Add tests for API throttling

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

Author: TG1999

vulnerabilities/tests/test_throttling.py

@@ -0,0 +1,48 @@
+#
+# Copyright (c) nexB Inc. and others. All rights reserved.
+# VulnerableCode is a trademark of nexB Inc.
+# SPDX-License-Identifier: Apache-2.0
+# See http://www.apache.org/licenses/LICENSE-2.0 for the license text.
+# See https://github.com/nexB/vulnerablecode for support or download.
+# See https://aboutcode.org for more information about nexB OSS projects.
+#
+
+from django.contrib.auth import get_user_model
+from rest_framework.test import APIClient
+from rest_framework.test import APITestCase
+
+User = get_user_model()
+
+
+class ThrottleApiTests(APITestCase):
+    def setUp(self):
+        # create a basic user
+        self.user = User.objects.create_user("username", "[email protected]", "secret")
+        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)
+
+        # create a staff user
+        self.staff_user = User.objects.create_user(
+            "staff", "[email protected]", "secret", is_staff=True
+        )
+        self.staff_auth = f"Token {self.staff_user.auth_token.key}"
+        self.staff_csrf_client = APIClient(enforce_csrf_checks=True)
+        self.staff_csrf_client.credentials(HTTP_AUTHORIZATION=self.staff_auth)
+
+    def test_api_throttling(self):
+
+        # A basic user can only access API 5 times a day
+        for i in range(0, 5):
+            response = self.csrf_client.get("/api/packages")
+            self.assertEqual(response.status_code, 200)
+            response = self.staff_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
+        self.assertEqual(response.status_code, 429)
+
+        response = self.staff_csrf_client.get("/api/packages", format="json")
+        # 200 - staff user can access API unlimited times
+        self.assertEqual(response.status_code, 200)

vulnerablecode/settings.py

@@ -150,9 +150,12 @@
 
 LOGIN_REDIRECT_URL = "/"
 LOGOUT_REDIRECT_URL = "/"
+THROTTLING_RATE = env.str("THROTTLING_RATE", default="1000/day")
 
 if IS_TESTS:
     VULNERABLECODEIO_REQUIRE_AUTHENTICATION = True
+    THROTTLING_RATE = "5/day"
+
 
 USE_L10N = True
 
@@ -187,7 +190,7 @@
     "DEFAULT_THROTTLE_CLASSES": [
         "vulnerabilities.throttling.ExceptionalUserRateThrottle",
     ],
-    "DEFAULT_THROTTLE_RATES": {"user": "1000/hour"},
+    "DEFAULT_THROTTLE_RATES": {"user": THROTTLING_RATE},
     "DEFAULT_PAGINATION_CLASS": "vulnerabilities.pagination.SmallResultSetPagination",
     # Limit the load on the Database returning a small number of records by default. https://github.com/nexB/vulnerablecode/issues/819
     "PAGE_SIZE": 10,