Add tests for user permission based API throttling

Signed-off-by: Keshav Priyadarshi <git@keshav.space>

Author: keshav-space

vulnerabilities/tests/test_api.py

@@ -11,6 +11,7 @@
 import os
 from urllib.parse import quote
 
+from django.core.cache import cache
 from django.test import TestCase
 from django.test import TransactionTestCase
 from django.test.client import RequestFactory
@@ -452,10 +453,12 @@ def add_aliases(vuln, aliases):
 
 class APIPerformanceTest(TestCase):
     def setUp(self):
-        self.user = ApiUser.objects.create_api_user(username="e@mail.com", is_staff=True)
-        self.auth = f"Token {self.user.auth_token.key}"
+        # Reset the api throttling to properly test the rate limit on anon users.
+        # DRF stores throttling state in cache, clear cache to reset throttling.
+        # See https://www.django-rest-framework.org/api-guide/throttling/#setting-up-the-cache
+        cache.clear()
+
         self.csrf_client = APIClient(enforce_csrf_checks=True)
-        self.csrf_client.credentials(HTTP_AUTHORIZATION=self.auth)
 
         # This setup creates the following data:
         # vulnerabilities: vul1, vul2, vul3
@@ -503,7 +506,7 @@ def setUp(self):
         set_as_fixing(package=self.pkg_2_13_2, vulnerability=self.vul1)
 
     def test_api_packages_all_num_queries(self):
-        with self.assertNumQueries(4):
+        with self.assertNumQueries(3):
             # There are 4 queries:
             # 1. SAVEPOINT
             # 2. Authenticating user
@@ -519,22 +522,22 @@ def test_api_packages_all_num_queries(self):
             ]
 
     def test_api_packages_single_num_queries(self):
-        with self.assertNumQueries(8):
+        with self.assertNumQueries(7):
             self.csrf_client.get(f"/api/packages/{self.pkg_2_14_0_rc1.id}", format="json")
 
     def test_api_packages_single_with_purl_in_query_num_queries(self):
-        with self.assertNumQueries(9):
+        with self.assertNumQueries(8):
             self.csrf_client.get(f"/api/packages/?purl={self.pkg_2_14_0_rc1.purl}", format="json")
 
     def test_api_packages_single_with_purl_no_version_in_query_num_queries(self):
-        with self.assertNumQueries(64):
+        with self.assertNumQueries(63):
             self.csrf_client.get(
                 f"/api/packages/?purl=pkg:maven/com.fasterxml.jackson.core/jackson-databind",
                 format="json",
             )
 
     def test_api_packages_bulk_search(self):
-        with self.assertNumQueries(45):
+        with self.assertNumQueries(44):
             packages = [self.pkg_2_12_6, self.pkg_2_12_6_1, self.pkg_2_13_1]
             purls = [p.purl for p in packages]
 
@@ -547,7 +550,7 @@ def test_api_packages_bulk_search(self):
             ).json()
 
     def test_api_packages_with_lookup(self):
-        with self.assertNumQueries(14):
+        with self.assertNumQueries(13):
             data = {"purl": self.pkg_2_12_6.purl}
 
             resp = self.csrf_client.post(
@@ -557,7 +560,7 @@ def test_api_packages_with_lookup(self):
             ).json()
 
     def test_api_packages_bulk_lookup(self):
-        with self.assertNumQueries(45):
+        with self.assertNumQueries(44):
             packages = [self.pkg_2_12_6, self.pkg_2_12_6_1, self.pkg_2_13_1]
             purls = [p.purl for p in packages]
 
@@ -572,10 +575,12 @@ def test_api_packages_bulk_lookup(self):
 
 class APITestCasePackage(TestCase):
     def setUp(self):
-        self.user = ApiUser.objects.create_api_user(username="e@mail.com", is_staff=True)
-        self.auth = f"Token {self.user.auth_token.key}"
+        # Reset the api throttling to properly test the rate limit on anon users.
+        # DRF stores throttling state in cache, clear cache to reset throttling.
+        # See https://www.django-rest-framework.org/api-guide/throttling/#setting-up-the-cache
+        cache.clear()
+
         self.csrf_client = APIClient(enforce_csrf_checks=True)
-        self.csrf_client.credentials(HTTP_AUTHORIZATION=self.auth)
 
         # This setup creates the following data:
         # vulnerabilities: vul1, vul2, vul3
@@ -766,7 +771,7 @@ def test_api_with_wrong_namespace_filter(self):
         self.assertEqual(response["count"], 0)
 
     def test_api_with_all_vulnerable_packages(self):
-        with self.assertNumQueries(4):
+        with self.assertNumQueries(3):
             # There are 4 queries:
             # 1. SAVEPOINT
             # 2. Authenticating user

vulnerabilities/tests/test_api_v2.py

@@ -61,10 +61,12 @@ def setUp(self):
         )
         self.reference2.vulnerabilities.add(self.vuln2)
 
-        self.user = ApiUser.objects.create_api_user(username="e@mail.com", is_staff=True)
-        self.auth = f"Token {self.user.auth_token.key}"
+        # Reset the api throttling to properly test the rate limit on anon users.
+        # DRF stores throttling state in cache, clear cache to reset throttling.
+        # See https://www.django-rest-framework.org/api-guide/throttling/#setting-up-the-cache
+        cache.clear()
+
         self.client = APIClient(enforce_csrf_checks=True)
-        self.client.credentials(HTTP_AUTHORIZATION=self.auth)
 
     def test_list_vulnerabilities(self):
         """
@@ -73,7 +75,7 @@ def test_list_vulnerabilities(self):
         """
         url = reverse("vulnerability-v2-list")
         response = self.client.get(url, format="json")
-        with self.assertNumQueries(5):
+        with self.assertNumQueries(4):
             response = self.client.get(url, format="json")
         self.assertEqual(response.status_code, status.HTTP_200_OK)
         self.assertIn("results", response.data)
@@ -88,7 +90,7 @@ def test_retrieve_vulnerability_detail(self):
         Test retrieving vulnerability details by vulnerability_id.
         """
         url = reverse("vulnerability-v2-detail", kwargs={"vulnerability_id": "VCID-1234"})
-        with self.assertNumQueries(8):
+        with self.assertNumQueries(7):
             response = self.client.get(url, format="json")
         self.assertEqual(response.status_code, status.HTTP_200_OK)
         self.assertEqual(response.data["vulnerability_id"], "VCID-1234")
@@ -102,7 +104,7 @@ def test_filter_vulnerability_by_vulnerability_id(self):
         Test filtering vulnerabilities by vulnerability_id.
         """
         url = reverse("vulnerability-v2-list")
-        with self.assertNumQueries(4):
+        with self.assertNumQueries(3):
             response = self.client.get(url, {"vulnerability_id": "VCID-1234"}, format="json")
         self.assertEqual(response.status_code, status.HTTP_200_OK)
         self.assertEqual(response.data["vulnerability_id"], "VCID-1234")
@@ -112,7 +114,7 @@ def test_filter_vulnerability_by_alias(self):
         Test filtering vulnerabilities by alias.
         """
         url = reverse("vulnerability-v2-list")
-        with self.assertNumQueries(5):
+        with self.assertNumQueries(4):
             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)
@@ -127,7 +129,7 @@ def test_filter_vulnerabilities_multiple_ids(self):
         Test filtering vulnerabilities by multiple vulnerability_ids.
         """
         url = reverse("vulnerability-v2-list")
-        with self.assertNumQueries(5):
+        with self.assertNumQueries(4):
             response = self.client.get(
                 url, {"vulnerability_id": ["VCID-1234", "VCID-5678"]}, format="json"
             )
@@ -139,7 +141,7 @@ def test_filter_vulnerabilities_multiple_aliases(self):
         Test filtering vulnerabilities by multiple aliases.
         """
         url = reverse("vulnerability-v2-list")
-        with self.assertNumQueries(5):
+        with self.assertNumQueries(4):
             response = self.client.get(
                 url, {"alias": ["CVE-2021-1234", "CVE-2021-5678"]}, format="json"
             )
@@ -152,7 +154,7 @@ def test_invalid_vulnerability_id(self):
         Should return 404 Not Found.
         """
         url = reverse("vulnerability-v2-detail", kwargs={"vulnerability_id": "VCID-9999"})
-        with self.assertNumQueries(5):
+        with self.assertNumQueries(4):
             response = self.client.get(url, format="json")
         self.assertEqual(response.status_code, status.HTTP_404_NOT_FOUND)
 
@@ -210,18 +212,20 @@ def setUp(self):
         self.package1.affected_by_vulnerabilities.add(self.vuln1)
         self.package2.fixing_vulnerabilities.add(self.vuln2)
 
-        self.user = ApiUser.objects.create_api_user(username="e@mail.com", is_staff=True)
-        self.auth = f"Token {self.user.auth_token.key}"
+        # Reset the api throttling to properly test the rate limit on anon users.
+        # DRF stores throttling state in cache, clear cache to reset throttling.
+        # See https://www.django-rest-framework.org/api-guide/throttling/#setting-up-the-cache
+        cache.clear()
+
         self.client = APIClient(enforce_csrf_checks=True)
-        self.client.credentials(HTTP_AUTHORIZATION=self.auth)
 
     def test_list_packages(self):
         """
         Test listing packages without filters.
         Should return a list of packages with their details and associated vulnerabilities.
         """
         url = reverse("package-v2-list")
-        with self.assertNumQueries(32):
+        with self.assertNumQueries(31):
             response = self.client.get(url, format="json")
         self.assertEqual(response.status_code, status.HTTP_200_OK)
         self.assertIn("results", response.data)
@@ -243,7 +247,7 @@ def test_filter_packages_by_purl(self):
         Test filtering packages by one or more PURLs.
         """
         url = reverse("package-v2-list")
-        with self.assertNumQueries(20):
+        with self.assertNumQueries(19):
             response = self.client.get(url, {"purl": "pkg:pypi/django@3.2"}, format="json")
         self.assertEqual(response.status_code, status.HTTP_200_OK)
         self.assertEqual(len(response.data["results"]["packages"]), 1)
@@ -254,7 +258,7 @@ def test_filter_packages_by_affected_vulnerability(self):
         Test filtering packages by affected_by_vulnerability.
         """
         url = reverse("package-v2-list")
-        with self.assertNumQueries(20):
+        with self.assertNumQueries(19):
             response = self.client.get(
                 url, {"affected_by_vulnerability": "VCID-1234"}, format="json"
             )
@@ -267,7 +271,7 @@ def test_filter_packages_by_fixing_vulnerability(self):
         Test filtering packages by fixing_vulnerability.
         """
         url = reverse("package-v2-list")
-        with self.assertNumQueries(18):
+        with self.assertNumQueries(17):
             response = self.client.get(url, {"fixing_vulnerability": "VCID-5678"}, format="json")
         self.assertEqual(response.status_code, status.HTTP_200_OK)
         self.assertEqual(len(response.data["results"]["packages"]), 1)
@@ -356,7 +360,7 @@ def test_invalid_vulnerability_filter(self):
         Should return an empty list.
         """
         url = reverse("package-v2-list")
-        with self.assertNumQueries(4):
+        with self.assertNumQueries(3):
             response = self.client.get(
                 url, {"affected_by_vulnerability": "VCID-9999"}, format="json"
             )
@@ -369,7 +373,7 @@ def test_invalid_purl_filter(self):
         Should return an empty list.
         """
         url = reverse("package-v2-list")
-        with self.assertNumQueries(4):
+        with self.assertNumQueries(3):
             response = self.client.get(
                 url, {"purl": "pkg:nonexistent/package@1.0.0"}, format="json"
             )
@@ -421,7 +425,7 @@ def test_bulk_lookup_with_valid_purls(self):
         """
         url = reverse("package-v2-bulk-lookup")
         data = {"purls": ["pkg:pypi/django@3.2", "pkg:npm/lodash@4.17.20"]}
-        with self.assertNumQueries(28):
+        with self.assertNumQueries(27):
             response = self.client.post(url, data, format="json")
         self.assertEqual(response.status_code, status.HTTP_200_OK)
         self.assertIn("packages", response.data)
@@ -446,7 +450,7 @@ def test_bulk_lookup_with_invalid_purls(self):
         """
         url = reverse("package-v2-bulk-lookup")
         data = {"purls": ["pkg:pypi/nonexistent@1.0.0", "pkg:npm/unknown@0.0.1"]}
-        with self.assertNumQueries(4):
+        with self.assertNumQueries(3):
             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
@@ -460,7 +464,7 @@ def test_bulk_lookup_with_empty_purls(self):
         """
         url = reverse("package-v2-bulk-lookup")
         data = {"purls": []}
-        with self.assertNumQueries(3):
+        with self.assertNumQueries(2):
             response = self.client.post(url, data, format="json")
         self.assertEqual(response.status_code, status.HTTP_400_BAD_REQUEST)
         self.assertIn("error", response.data)
@@ -474,7 +478,7 @@ def test_bulk_search_with_valid_purls(self):
         """
         url = reverse("package-v2-bulk-search")
         data = {"purls": ["pkg:pypi/django@3.2", "pkg:npm/lodash@4.17.20"]}
-        with self.assertNumQueries(28):
+        with self.assertNumQueries(27):
             response = self.client.post(url, data, format="json")
         self.assertEqual(response.status_code, status.HTTP_200_OK)
         self.assertIn("packages", response.data)
@@ -502,7 +506,7 @@ def test_bulk_search_with_purl_only_true(self):
             "purls": ["pkg:pypi/django@3.2", "pkg:npm/lodash@4.17.20"],
             "purl_only": True,
         }
-        with self.assertNumQueries(17):
+        with self.assertNumQueries(16):
             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
@@ -529,7 +533,7 @@ def test_bulk_search_with_plain_purl_true(self):
             "purls": ["pkg:pypi/django@3.2", "pkg:pypi/django@3.2?extension=tar.gz"],
             "plain_purl": True,
         }
-        with self.assertNumQueries(16):
+        with self.assertNumQueries(15):
             response = self.client.post(url, data, format="json")
         self.assertEqual(response.status_code, status.HTTP_200_OK)
         self.assertIn("packages", response.data)
@@ -550,7 +554,7 @@ def test_bulk_search_with_purl_only_and_plain_purl_true(self):
             "purl_only": True,
             "plain_purl": True,
         }
-        with self.assertNumQueries(11):
+        with self.assertNumQueries(10):
             response = self.client.post(url, data, format="json")
         self.assertEqual(response.status_code, status.HTTP_200_OK)
         # Response should be a list of plain PURLs
@@ -566,7 +570,7 @@ def test_bulk_search_with_invalid_purls(self):
         """
         url = reverse("package-v2-bulk-search")
         data = {"purls": ["pkg:pypi/nonexistent@1.0.0", "pkg:npm/unknown@0.0.1"]}
-        with self.assertNumQueries(4):
+        with self.assertNumQueries(3):
             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
@@ -580,7 +584,7 @@ def test_bulk_search_with_empty_purls(self):
         """
         url = reverse("package-v2-bulk-search")
         data = {"purls": []}
-        with self.assertNumQueries(3):
+        with self.assertNumQueries(2):
             response = self.client.post(url, data, format="json")
         self.assertEqual(response.status_code, status.HTTP_400_BAD_REQUEST)
         self.assertIn("error", response.data)
@@ -592,7 +596,7 @@ def test_all_vulnerable_packages(self):
         Test the 'all' endpoint that returns all vulnerable package URLs.
         """
         url = reverse("package-v2-all")
-        with self.assertNumQueries(4):
+        with self.assertNumQueries(3):
             response = self.client.get(url, format="json")
         self.assertEqual(response.status_code, status.HTTP_200_OK)
         # Since package1 is vulnerable, it should be returned
@@ -606,7 +610,7 @@ def test_lookup_with_valid_purl(self):
         """
         url = reverse("package-v2-lookup")
         data = {"purl": "pkg:pypi/django@3.2"}
-        with self.assertNumQueries(13):
+        with self.assertNumQueries(12):
             response = self.client.post(url, data, format="json")
         self.assertEqual(response.status_code, status.HTTP_200_OK)
         self.assertEqual(1, len(response.data))
@@ -635,7 +639,7 @@ def test_lookup_with_invalid_purl(self):
         """
         url = reverse("package-v2-lookup")
         data = {"purl": "pkg:pypi/nonexistent@1.0.0"}
-        with self.assertNumQueries(4):
+        with self.assertNumQueries(3):
             response = self.client.post(url, data, format="json")
         self.assertEqual(response.status_code, status.HTTP_200_OK)
         # No packages or vulnerabilities should be returned
@@ -648,7 +652,7 @@ def test_lookup_with_missing_purl(self):
         """
         url = reverse("package-v2-lookup")
         data = {}
-        with self.assertNumQueries(3):
+        with self.assertNumQueries(2):
             response = self.client.post(url, data, format="json")
         self.assertEqual(response.status_code, status.HTTP_400_BAD_REQUEST)
         self.assertIn("error", response.data)
@@ -662,7 +666,7 @@ def test_lookup_with_invalid_purl_format(self):
         """
         url = reverse("package-v2-lookup")
         data = {"purl": "invalid_purl_format"}
-        with self.assertNumQueries(4):
+        with self.assertNumQueries(3):
             response = self.client.post(url, data, format="json")
         self.assertEqual(response.status_code, status.HTTP_200_OK)
         # No packages or vulnerabilities should be returned