|
| 1 | +from django.db import connection |
| 2 | +from django.test.utils import CaptureQueriesContext |
| 3 | +from django.utils.timezone import now |
| 4 | +from rest_framework.test import APITestCase |
| 5 | + |
| 6 | +from dojo.models import ( |
| 7 | + Dojo_User, |
| 8 | + Engagement, |
| 9 | + Finding, |
| 10 | + NoteHistory, |
| 11 | + Notes, |
| 12 | + Product, |
| 13 | + Product_Type, |
| 14 | + Risk_Acceptance, |
| 15 | + Test, |
| 16 | + Test_Type, |
| 17 | +) |
| 18 | + |
| 19 | + |
| 20 | +class TestNotesListNPlusOne(APITestCase): |
| 21 | + |
| 22 | + """ |
| 23 | + Regression: list endpoints that embed NoteSerializer (findings, engagements, tests, |
| 24 | + risk_acceptance) must load note authors/editors/history in bulk via notes_prefetch(). |
| 25 | + A flat prefetch of the notes relation leaves NoteSerializer to lazy-load author, editor, |
| 26 | + note_type and history__current_editor per note, so the query count grows with every note. |
| 27 | + Each test pins the fix by asserting the query count is identical with 1 and with 5 notes. |
| 28 | + """ |
| 29 | + |
| 30 | + @classmethod |
| 31 | + def setUpTestData(cls): |
| 32 | + cls.user = Dojo_User.objects.create(username="np1_user", is_staff=True, is_superuser=True) |
| 33 | + cls.prod_type = Product_Type.objects.create(name="NP1 Product Type") |
| 34 | + cls.product = Product.objects.create(name="NP1 Product", prod_type=cls.prod_type, description="NP1") |
| 35 | + cls.engagement = Engagement.objects.create( |
| 36 | + name="NP1 Engagement", product=cls.product, target_start=now(), target_end=now(), |
| 37 | + ) |
| 38 | + cls.test_type = Test_Type.objects.create(name="NP1 Test Type") |
| 39 | + cls.test = Test.objects.create( |
| 40 | + title="NP1 Test", engagement=cls.engagement, test_type=cls.test_type, |
| 41 | + target_start=now(), target_end=now(), |
| 42 | + ) |
| 43 | + cls.finding = Finding.objects.create(title="NP1 Finding", test=cls.test, reporter=cls.user, severity="High") |
| 44 | + cls.risk_acceptance = Risk_Acceptance.objects.create(name="NP1 RA", owner=cls.user) |
| 45 | + cls.risk_acceptance.engagement_set.add(cls.engagement) |
| 46 | + |
| 47 | + def setUp(self): |
| 48 | + self.client.force_authenticate(user=self.user) |
| 49 | + |
| 50 | + def _add_note(self, obj): |
| 51 | + # An edited note with a history entry exercises every relation NoteSerializer renders. |
| 52 | + note = Notes.objects.create(entry="np1 note", author=self.user, edited=True, editor=self.user) |
| 53 | + history = NoteHistory.objects.create(data="np1 history", current_editor=self.user) |
| 54 | + note.history.add(history) |
| 55 | + obj.notes.add(note) |
| 56 | + |
| 57 | + def _query_count(self, url): |
| 58 | + with CaptureQueriesContext(connection) as ctx: |
| 59 | + response = self.client.get(url) |
| 60 | + self.assertEqual(response.status_code, 200, response.content[:1000]) |
| 61 | + return len(ctx.captured_queries) |
| 62 | + |
| 63 | + def _assert_count_independent_of_notes(self, obj, url): |
| 64 | + self._add_note(obj) |
| 65 | + self._query_count(url) # warm-up: fills ContentType and other first-request caches |
| 66 | + with_one_note = self._query_count(url) |
| 67 | + for _ in range(4): |
| 68 | + self._add_note(obj) |
| 69 | + with_five_notes = self._query_count(url) |
| 70 | + self.assertEqual( |
| 71 | + with_one_note, with_five_notes, |
| 72 | + f"{url} ran {with_five_notes - with_one_note} extra queries for 4 extra notes (N+1 on notes)", |
| 73 | + ) |
| 74 | + |
| 75 | + def test_finding_list_query_count_independent_of_notes(self): |
| 76 | + self._assert_count_independent_of_notes(self.finding, f"/api/v2/findings/?id={self.finding.id}") |
| 77 | + |
| 78 | + def test_engagement_list_query_count_independent_of_notes(self): |
| 79 | + self._assert_count_independent_of_notes(self.engagement, f"/api/v2/engagements/?id={self.engagement.id}") |
| 80 | + |
| 81 | + def test_test_list_query_count_independent_of_notes(self): |
| 82 | + self._assert_count_independent_of_notes(self.test, f"/api/v2/tests/?id={self.test.id}") |
| 83 | + |
| 84 | + def test_risk_acceptance_list_query_count_independent_of_notes(self): |
| 85 | + self._assert_count_independent_of_notes(self.risk_acceptance, f"/api/v2/risk_acceptance/?id={self.risk_acceptance.id}") |
0 commit comments