Add test for group based throttling

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

Author: keshav-space

vulnerabilities/tests/test_api.py

@@ -452,7 +452,7 @@ def add_aliases(vuln, aliases):
 
 class APIPerformanceTest(TestCase):
     def setUp(self):
-        self.user = ApiUser.objects.create_api_user(username="e@mail.com")
+        self.user = ApiUser.objects.create_api_user(username="e@mail.com", is_staff=True)
         self.auth = f"Token {self.user.auth_token.key}"
         self.csrf_client = APIClient(enforce_csrf_checks=True)
         self.csrf_client.credentials(HTTP_AUTHORIZATION=self.auth)
@@ -572,7 +572,7 @@ def test_api_packages_bulk_lookup(self):
 
 class APITestCasePackage(TestCase):
     def setUp(self):
-        self.user = ApiUser.objects.create_api_user(username="e@mail.com")
+        self.user = ApiUser.objects.create_api_user(username="e@mail.com", is_staff=True)
         self.auth = f"Token {self.user.auth_token.key}"
         self.csrf_client = APIClient(enforce_csrf_checks=True)
         self.csrf_client.credentials(HTTP_AUTHORIZATION=self.auth)

vulnerabilities/tests/test_api_v2.py

@@ -61,7 +61,7 @@ def setUp(self):
         )
         self.reference2.vulnerabilities.add(self.vuln2)
 
-        self.user = ApiUser.objects.create_api_user(username="e@mail.com")
+        self.user = ApiUser.objects.create_api_user(username="e@mail.com", is_staff=True)
         self.auth = f"Token {self.user.auth_token.key}"
         self.client = APIClient(enforce_csrf_checks=True)
         self.client.credentials(HTTP_AUTHORIZATION=self.auth)
@@ -210,7 +210,7 @@ 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")
+        self.user = ApiUser.objects.create_api_user(username="e@mail.com", is_staff=True)
         self.auth = f"Token {self.user.auth_token.key}"
         self.client = APIClient(enforce_csrf_checks=True)
         self.client.credentials(HTTP_AUTHORIZATION=self.auth)

vulnerabilities/tests/test_throttling.py

@@ -9,25 +9,43 @@
 
 import json
 
+from django.contrib.auth.models import Group
 from django.core.cache import cache
 from rest_framework.test import APIClient
 from rest_framework.test import APITestCase
 
 from vulnerabilities.models import ApiUser
 
 
-class ThrottleApiTests(APITestCase):
+class GroupUserRateThrottleApiTests(APITestCase):
     def setUp(self):
         # 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()
 
-        # create a basic user
-        self.user = ApiUser.objects.create_api_user(username="e@mail.com")
-        self.auth = f"Token {self.user.auth_token.key}"
-        self.csrf_client = APIClient(enforce_csrf_checks=True)
-        self.csrf_client.credentials(HTTP_AUTHORIZATION=self.auth)
+        # User in bronze group
+        self.bronze_user = ApiUser.objects.create_api_user(username="bronze@mail.com")
+        bronze, _ = Group.objects.get_or_create(name="bronze")
+        self.bronze_user.groups.clear()
+        self.bronze_user.groups.add(bronze)
+        self.bronze_auth = f"Token {self.bronze_user.auth_token.key}"
+        self.bronze_user_csrf_client = APIClient(enforce_csrf_checks=True)
+        self.bronze_user_csrf_client.credentials(HTTP_AUTHORIZATION=self.bronze_auth)
+
+        # User in silver group (default group for api user)
+        self.silver_user = ApiUser.objects.create_api_user(username="silver@mail.com")
+        self.silver_auth = f"Token {self.silver_user.auth_token.key}"
+        self.silver_user_csrf_client = APIClient(enforce_csrf_checks=True)
+        self.silver_user_csrf_client.credentials(HTTP_AUTHORIZATION=self.silver_auth)
+
+        # User in gold group
+        self.gold_user = ApiUser.objects.create_api_user(username="gold@mail.com")
+        gold, _ = Group.objects.get_or_create(name="gold")
+        self.gold_user.groups.add(gold)
+        self.gold_auth = f"Token {self.gold_user.auth_token.key}"
+        self.gold_user_csrf_client = APIClient(enforce_csrf_checks=True)
+        self.gold_user_csrf_client.credentials(HTTP_AUTHORIZATION=self.gold_auth)
 
         # create a staff user
         self.staff_user = ApiUser.objects.create_api_user(username="staff@mail.com", is_staff=True)
@@ -39,16 +57,34 @@ def setUp(self):
         self.csrf_client_anon_1 = APIClient(enforce_csrf_checks=True)
 
     def test_package_endpoint_throttling(self):
-        for i in range(0, 20):
-            response = self.csrf_client.get("/api/packages")
+        for i in range(0, 15):
+            response = self.bronze_user_csrf_client.get("/api/packages")
             self.assertEqual(response.status_code, 200)
-            response = self.staff_csrf_client.get("/api/packages")
+
+        response = self.bronze_user_csrf_client.get("/api/packages")
+        # 429 - too many requests for bronze user
+        self.assertEqual(response.status_code, 429)
+
+        for i in range(0, 20):
+            response = self.silver_user_csrf_client.get("/api/packages")
             self.assertEqual(response.status_code, 200)
 
-        response = self.csrf_client.get("/api/packages")
-        # 429 - too many requests for basic user
+        response = self.silver_user_csrf_client.get("/api/packages")
+        # 429 - too many requests for silver user
         self.assertEqual(response.status_code, 429)
 
+        for i in range(0, 30):
+            response = self.gold_user_csrf_client.get("/api/packages")
+            self.assertEqual(response.status_code, 200)
+
+        response = self.gold_user_csrf_client.get("/api/packages", format="json")
+        # 200 - gold user can access API unlimited times
+        self.assertEqual(response.status_code, 200)
+
+        for i in range(0, 30):
+            response = self.staff_csrf_client.get("/api/packages")
+            self.assertEqual(response.status_code, 200)
+
         response = self.staff_csrf_client.get("/api/packages", format="json")
         # 200 - staff user can access API unlimited times
         self.assertEqual(response.status_code, 200)