Address review comments

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

Author: TG1999

vulnerabilities/tests/test_fix_api.py

@@ -259,13 +259,12 @@ def test_api_with_single_vulnerability_and_vulnerable_package(self):
         }
 
     def test_api_with_all_vulnerable_packages(self):
-        with self.assertNumQueries(5):
+        with self.assertNumQueries(4):
             # There are 4 queries:
             # 1. SAVEPOINT
             # 2. Authenticating user
-            # 3. Checking if user is staff user for throttling purposes
-            # 4. Get all vulnerable packages
-            # 5. RELEASE SAVEPOINT
+            # 3. Get all vulnerable packages
+            # 4. RELEASE SAVEPOINT
             response = self.csrf_client.get(f"/api/packages/all", format="json").data
             assert len(response) == 11
             assert response == [

vulnerabilities/throttling.py

@@ -13,40 +13,12 @@
 User = get_user_model()
 
 
-class ExceptionalUserRateThrottle(UserRateThrottle):
+class StaffUserRateThrottle(UserRateThrottle):
     def allow_request(self, request, view):
         """
-        Give special access to a few special accounts.
-
-        Mirrors code in super class with minor tweaks.
+        Do not apply throttling for superusers and admins.
         """
-        if self.rate is None:
-            return True
-
-        self.key = self.get_cache_key(request, view)
-        if self.key is None:
+        if request.user.is_superuser or request.user.is_staff:
             return True
 
-        self.history = self.cache.get(self.key, [])
-        self.now = self.timer()
-
-        # Adjust if user has special privileges.
-
-        user = User.objects.get(username=request.user.username)
-
-        if user:
-            if user.is_superuser or user.is_staff:
-                # No throttling for superusers or staff.
-                return True
-
-            else:
-                self.num_requests = self.num_requests
-                self.duration = self.duration
-
-        # Drop any requests from the history which have now passed the
-        # throttle duration
-        while self.history and self.history[-1] <= self.now - self.duration:
-            self.history.pop()
-        if len(self.history) >= self.num_requests:
-            return self.throttle_failure()
-        return self.throttle_success()
+        return super().allow_request(request, view)

vulnerablecode/settings.py

@@ -188,7 +188,7 @@
         "rest_framework.filters.SearchFilter",
     ),
     "DEFAULT_THROTTLE_CLASSES": [
-        "vulnerabilities.throttling.ExceptionalUserRateThrottle",
+        "vulnerabilities.throttling.StaffUserRateThrottle",
     ],
     "DEFAULT_THROTTLE_RATES": {"user": THROTTLING_RATE},
     "DEFAULT_PAGINATION_CLASS": "vulnerabilities.pagination.SmallResultSetPagination",