Add tests for num queries for views (#1730)

* Add tests for num queries for views

Signed-off-by: Tushar Goel <tushar.goel.dav@gmail.com>

* Add tests for num queries for views

Signed-off-by: Tushar Goel <tushar.goel.dav@gmail.com>

* Add tests for num queries for views

Signed-off-by: Tushar Goel <tushar.goel.dav@gmail.com>

---------

Signed-off-by: Tushar Goel <tushar.goel.dav@gmail.com>

Author: TG1999

vulnerabilities/tests/test_view.py

@@ -8,16 +8,21 @@
 #
 
 import os
+import time
 
 import pytest
 from django.test import Client
 from django.test import TestCase
 from packageurl import PackageURL
 from univers import versions
 
+from vulnerabilities import models
+from vulnerabilities.models import AffectedByPackageRelatedVulnerability
 from vulnerabilities.models import Alias
+from vulnerabilities.models import FixingPackageRelatedVulnerability
 from vulnerabilities.models import Package
 from vulnerabilities.models import Vulnerability
+from vulnerabilities.models import VulnerabilitySeverity
 from vulnerabilities.templatetags.url_filters import url_quote_filter
 from vulnerabilities.views import PackageDetails
 from vulnerabilities.views import PackageSearch
@@ -273,3 +278,56 @@ class TestCustomFilters:
     def test_url_quote_filter(self, input_value, expected_output):
         filtered = url_quote_filter(input_value)
         assert filtered == expected_output
+
+
+class VulnerabilitySearchTestCaseWithPackages(TestCase):
+    def setUp(self):
+        self.vuln1 = Vulnerability.objects.create(vulnerability_id="VCID-1", summary="Vuln 1")
+        self.vuln2 = Vulnerability.objects.create(vulnerability_id="VCID-2", summary="Vuln 2")
+        self.vuln3 = Vulnerability.objects.create(vulnerability_id="VCID-3", summary="Vuln 3")
+        self.vuln4 = Vulnerability.objects.create(vulnerability_id="VCID-4", summary="Vuln 4")
+        self.vuln5 = Vulnerability.objects.create(vulnerability_id="VCID-5", summary="Vuln 5")
+
+        self.package1 = Package.objects.create(type="pypi", name="django", version="1.0.0")
+        self.package2 = Package.objects.create(type="pypi", name="django", version="2.0.0")
+        self.package3 = Package.objects.create(type="pypi", name="django", version="3.0.0")
+
+        AffectedByPackageRelatedVulnerability.objects.create(
+            package=self.package1, vulnerability=self.vuln1
+        )
+        AffectedByPackageRelatedVulnerability.objects.create(
+            package=self.package1, vulnerability=self.vuln2
+        )
+        AffectedByPackageRelatedVulnerability.objects.create(
+            package=self.package2, vulnerability=self.vuln3
+        )
+        AffectedByPackageRelatedVulnerability.objects.create(
+            package=self.package2, vulnerability=self.vuln4
+        )
+
+        FixingPackageRelatedVulnerability.objects.create(
+            package=self.package3, vulnerability=self.vuln5
+        )
+
+        self.severity1 = VulnerabilitySeverity.objects.create(
+            scoring_system="CVSSv3",
+            value="9.8",
+            scoring_elements="AV:N/AC:L/PR:N/UI:N/S:U/C:H/I:H/A:H",
+        )
+        self.severity2 = VulnerabilitySeverity.objects.create(
+            scoring_system="CVSSv3",
+            value="7.5",
+            scoring_elements="AV:N/AC:L/PR:N/UI:N/S:U/C:H/I:H/A:H",
+        )
+
+        self.vuln1.severities.add(self.severity1)
+        self.vuln1.severities.add(self.severity2)
+        self.vuln1.save()
+
+    def test_aggregate_fixed_and_affected_packages(self):
+        with self.assertNumQueries(11):
+            start_time = time.time()
+            response = self.client.get(f"/vulnerabilities/{self.vuln1.vulnerability_id}")
+            end_time = time.time()
+            assert end_time - start_time < 0.05
+            self.assertEqual(response.status_code, 200)