Update tests

Signed-off-by: Tushar Goel <[email protected]>

Author: TG1999

vulnerabilities/tests/test_api_v2.py

@@ -54,6 +54,7 @@ def setUp(self):
             url="https://example.com/ref2", reference_type="exploit", reference_id="REF-2"
         )
         self.reference2.vulnerabilities.add(self.vuln2)
+
         self.user = ApiUser.objects.create_api_user(username="[email protected]")
         self.auth = f"Token {self.user.auth_token.key}"
         self.client = APIClient(enforce_csrf_checks=True)
@@ -62,13 +63,16 @@ def setUp(self):
     def test_list_vulnerabilities(self):
         """
         Test listing vulnerabilities without filters.
-        Should return a list of vulnerabilities with IDs and URLs.
+        Should return a paginated response with vulnerabilities dictionary.
         """
         url = reverse("vulnerability-v2-list")
         response = self.client.get(url, format="json")
         self.assertEqual(response.status_code, status.HTTP_200_OK)
+        self.assertIn("results", response.data)
         self.assertIn("vulnerabilities", response.data["results"])
         self.assertEqual(len(response.data["results"]["vulnerabilities"]), 2)
+        self.assertIn("VCID-1234", response.data["results"]["vulnerabilities"])
+        self.assertIn("VCID-5678", response.data["results"]["vulnerabilities"])
         self.assertTrue("url" in response.data["results"]["vulnerabilities"]["VCID-1234"])
 
     def test_retrieve_vulnerability_detail(self):
@@ -100,6 +104,8 @@ def test_filter_vulnerability_by_alias(self):
         url = reverse("vulnerability-v2-list")
         response = self.client.get(url, {"alias": "CVE-2021-5678"}, format="json")
         self.assertEqual(response.status_code, status.HTTP_200_OK)
+        self.assertIn("results", response.data)
+        self.assertIn("vulnerabilities", response.data["results"])
         self.assertEqual(
             response.data["results"]["vulnerabilities"]["VCID-5678"]["vulnerability_id"],
             "VCID-5678",
@@ -159,9 +165,13 @@ def test_list_vulnerabilities_pagination(self):
         response = self.client.get(url, format="json")
         self.assertEqual(response.status_code, status.HTTP_200_OK)
         self.assertIn("results", response.data)
+        self.assertIn("vulnerabilities", response.data["results"])
         self.assertIn("next", response.data)
         self.assertIn("previous", response.data)
-        self.assertEqual(len(response.data["results"]), 1)  # Assuming default page size is 10
+        # The 'vulnerabilities' dictionary should contain vulnerabilities up to the page limit
+        self.assertEqual(
+            len(response.data["results"]["vulnerabilities"]), 10
+        )  # Assuming default page size is 10
 
 
 class PackageV2ViewSetTest(APITestCase):
@@ -185,6 +195,7 @@ def setUp(self):
         # Associate packages with vulnerabilities
         self.package1.affected_by_vulnerabilities.add(self.vuln1)
         self.package2.fixing_vulnerabilities.add(self.vuln2)
+
         self.user = ApiUser.objects.create_api_user(username="[email protected]")
         self.auth = f"Token {self.user.auth_token.key}"
         self.client = APIClient(enforce_csrf_checks=True)
@@ -193,13 +204,24 @@ def setUp(self):
     def test_list_packages(self):
         """
         Test listing packages without filters.
-        Should return a list of packages with their details.
+        Should return a list of packages with their details and associated vulnerabilities.
         """
         url = reverse("package-v2-list")
         response = self.client.get(url, format="json")
         self.assertEqual(response.status_code, status.HTTP_200_OK)
+        self.assertIn("results", response.data)
         self.assertIn("packages", response.data["results"])
+        self.assertIn("vulnerabilities", response.data["results"])
         self.assertEqual(len(response.data["results"]["packages"]), 2)
+        # Verify that vulnerabilities are included
+        self.assertIsInstance(response.data["results"]["vulnerabilities"], dict)
+        package_vulns = set()
+        for package in response.data["results"]["packages"]:
+            package_vulns.update(package["affected_by_vulnerabilities"])
+            package_vulns.update(package["fixing_vulnerabilities"])
+        self.assertTrue(
+            all(vuln_id in response.data["results"]["vulnerabilities"] for vuln_id in package_vulns)
+        )
 
     def test_filter_packages_by_purl(self):
         """
@@ -264,9 +286,13 @@ def test_list_packages_pagination(self):
         response = self.client.get(url, format="json")
         self.assertEqual(response.status_code, status.HTTP_200_OK)
         self.assertIn("results", response.data)
+        self.assertIn("packages", response.data["results"])
+        self.assertIn("vulnerabilities", response.data["results"])
         self.assertIn("next", response.data)
         self.assertIn("previous", response.data)
-        self.assertEqual(len(response.data["results"]), 1)  # Assuming default page size is 10
+        self.assertEqual(
+            len(response.data["results"]["packages"]), 10
+        )  # Assuming default page size is 10
 
     def test_invalid_vulnerability_filter(self):
         """
@@ -309,16 +335,27 @@ def test_get_fixing_vulnerabilities(self):
     def test_bulk_lookup_with_valid_purls(self):
         """
         Test bulk lookup with valid PURLs.
+        Should return packages and their associated vulnerabilities.
         """
         url = reverse("package-v2-bulk-lookup")
         data = {"purls": ["pkg:pypi/[email protected]", "pkg:npm/[email protected]"]}
         response = self.client.post(url, data, format="json")
         self.assertEqual(response.status_code, status.HTTP_200_OK)
-        self.assertEqual(len(response.data), 2)
+        self.assertIn("packages", response.data)
+        self.assertIn("vulnerabilities", response.data)
+        self.assertEqual(len(response.data["packages"]), 2)
         # Verify that the returned data matches the packages
-        purls = [package["purl"] for package in response.data]
+        purls = [package["purl"] for package in response.data["packages"]]
         self.assertIn("pkg:pypi/[email protected]", purls)
         self.assertIn("pkg:npm/[email protected]", purls)
+        # Verify that vulnerabilities are included
+        package_vulns = set()
+        for package in response.data["packages"]:
+            package_vulns.update(package["affected_by_vulnerabilities"])
+            package_vulns.update(package["fixing_vulnerabilities"])
+        self.assertTrue(
+            all(vuln_id in response.data["vulnerabilities"] for vuln_id in package_vulns)
+        )
 
     def test_bulk_lookup_with_invalid_purls(self):
         """
@@ -329,7 +366,8 @@ def test_bulk_lookup_with_invalid_purls(self):
         response = self.client.post(url, data, format="json")
         self.assertEqual(response.status_code, status.HTTP_200_OK)
         # Since the packages don't exist, the response should be empty
-        self.assertEqual(len(response.data), 0)
+        self.assertEqual(len(response.data["packages"]), 0)
+        self.assertEqual(len(response.data["vulnerabilities"]), 0)
 
     def test_bulk_lookup_with_empty_purls(self):
         """
@@ -347,23 +385,37 @@ def test_bulk_lookup_with_empty_purls(self):
     def test_bulk_search_with_valid_purls(self):
         """
         Test bulk search with valid PURLs.
+        Should return packages and their associated vulnerabilities.
         """
         url = reverse("package-v2-bulk-search")
         data = {"purls": ["pkg:pypi/[email protected]", "pkg:npm/[email protected]"]}
         response = self.client.post(url, data, format="json")
         self.assertEqual(response.status_code, status.HTTP_200_OK)
-        self.assertEqual(len(response.data), 2)
-        purls = [package["purl"] for package in response.data]
+        self.assertIn("packages", response.data)
+        self.assertIn("vulnerabilities", response.data)
+        self.assertEqual(len(response.data["packages"]), 2)
+        purls = [package["purl"] for package in response.data["packages"]]
         self.assertIn("pkg:pypi/[email protected]", purls)
         self.assertIn("pkg:npm/[email protected]", purls)
+        # Verify that vulnerabilities are included
+        package_vulns = set()
+        for package in response.data["packages"]:
+            package_vulns.update(package["affected_by_vulnerabilities"])
+            package_vulns.update(package["fixing_vulnerabilities"])
+        self.assertTrue(
+            all(vuln_id in response.data["vulnerabilities"] for vuln_id in package_vulns)
+        )
 
     def test_bulk_search_with_purl_only_true(self):
         """
         Test bulk search with purl_only set to True.
         Should return only the PURLs of vulnerable packages.
         """
         url = reverse("package-v2-bulk-search")
-        data = {"purls": ["pkg:pypi/[email protected]", "pkg:npm/[email protected]"], "purl_only": True}
+        data = {
+            "purls": ["pkg:pypi/[email protected]", "pkg:npm/[email protected]"],
+            "purl_only": True,
+        }
         response = self.client.post(url, data, format="json")
         self.assertEqual(response.status_code, status.HTTP_200_OK)
         # Since purl_only=True, response should be a list of PURLs
@@ -375,15 +427,29 @@ def test_bulk_search_with_purl_only_true(self):
     def test_bulk_search_with_plain_purl_true(self):
         """
         Test bulk search with plain_purl set to True.
-        """
+        Should return packages grouped by plain PURLs.
+        """
+        # Create another package with the same name and version but different qualifiers
+        Package.objects.create(
+            name="django",
+            version="3.2",
+            type="pypi",
+            qualifiers={"extension": "tar.gz"},
+        )
+
         url = reverse("package-v2-bulk-search")
-        data = {"purls": ["pkg:pypi/[email protected]", "pkg:pypi/[email protected]"], "plain_purl": True}
+        data = {
+            "purls": ["pkg:pypi/[email protected]", "pkg:pypi/[email protected]?extension=tar.gz"],
+            "plain_purl": True,
+        }
         response = self.client.post(url, data, format="json")
         self.assertEqual(response.status_code, status.HTTP_200_OK)
-        # Since plain_purl=True, packages with the same name and version are grouped
-        self.assertEqual(len(response.data), 1)
-        purls = [package["purl"] for package in response.data]
-        self.assertIn("pkg:pypi/[email protected]", purls[0] or "pkg:pypi/[email protected]" in purls[0])
+        self.assertIn("packages", response.data)
+        self.assertIn("vulnerabilities", response.data)
+        # Since plain_purl=True, packages with the same type, namespace, name, version are grouped
+        self.assertEqual(len(response.data["packages"]), 1)
+        purl = response.data["packages"][0]["purl"]
+        self.assertTrue(purl.startswith("pkg:pypi/[email protected]"))
 
     def test_bulk_search_with_purl_only_and_plain_purl_true(self):
         """
@@ -407,12 +473,15 @@ def test_bulk_search_with_purl_only_and_plain_purl_true(self):
     def test_bulk_search_with_invalid_purls(self):
         """
         Test bulk search with invalid PURLs.
+        Should return an empty response.
         """
         url = reverse("package-v2-bulk-search")
         data = {"purls": ["pkg:pypi/[email protected]", "pkg:npm/[email protected]"]}
         response = self.client.post(url, data, format="json")
         self.assertEqual(response.status_code, status.HTTP_200_OK)
-        self.assertEqual(len(response.data), 0)
+        # Since the packages don't exist, the response should be empty
+        self.assertEqual(len(response.data["packages"]), 0)
+        self.assertEqual(len(response.data["vulnerabilities"]), 0)
 
     def test_bulk_search_with_empty_purls(self):
         """
@@ -434,32 +503,39 @@ def test_all_vulnerable_packages(self):
         url = reverse("package-v2-all")
         response = self.client.get(url, format="json")
         self.assertEqual(response.status_code, status.HTTP_200_OK)
-        # Since package1 and package3 are vulnerable, they should be returned
+        # Since package1 is vulnerable, it should be returned
         expected_purls = ["pkg:pypi/[email protected]"]
         self.assertEqual(sorted(response.data), sorted(expected_purls))
 
     def test_lookup_with_valid_purl(self):
         """
         Test the 'lookup' endpoint with a valid PURL.
+        Should return the package and its associated vulnerabilities.
         """
         url = reverse("package-v2-lookup")
         data = {"purl": "pkg:pypi/[email protected]"}
         response = self.client.post(url, data, format="json")
         self.assertEqual(response.status_code, status.HTTP_200_OK)
-        self.assertEqual(len(response.data), 1)
+        self.assertEqual(1, len(response.data))
+        self.assertIn("purl", response.data[0])
+        self.assertIn("affected_by_vulnerabilities", response.data[0])
+        self.assertIn("fixing_vulnerabilities", response.data[0])
+        self.assertIn("next_non_vulnerable_version", response.data[0])
+        self.assertIn("latest_non_vulnerable_version", response.data[0])
         self.assertEqual(response.data[0]["purl"], "pkg:pypi/[email protected]")
         self.assertEqual(response.data[0]["affected_by_vulnerabilities"], ["VCID-1234"])
+        self.assertEqual(response.data[0]["fixing_vulnerabilities"], [])
 
     def test_lookup_with_invalid_purl(self):
         """
         Test the 'lookup' endpoint with a PURL that does not exist.
-        Should return an empty list.
+        Should return empty packages and vulnerabilities.
         """
         url = reverse("package-v2-lookup")
         data = {"purl": "pkg:pypi/[email protected]"}
         response = self.client.post(url, data, format="json")
         self.assertEqual(response.status_code, status.HTTP_200_OK)
-        # No packages should be returned
+        # No packages or vulnerabilities should be returned
         self.assertEqual(len(response.data), 0)
 
     def test_lookup_with_missing_purl(self):
@@ -478,9 +554,11 @@ def test_lookup_with_missing_purl(self):
     def test_lookup_with_invalid_purl_format(self):
         """
         Test the 'lookup' endpoint with an invalid PURL format.
-        Should return 400 Bad Request.
+        Should return empty packages and vulnerabilities.
         """
         url = reverse("package-v2-lookup")
         data = {"purl": "invalid_purl_format"}
         response = self.client.post(url, data, format="json")
         self.assertEqual(response.status_code, status.HTTP_200_OK)
+        # No packages or vulnerabilities should be returned
+        self.assertEqual(len(response.data), 0)