From 7bb2e6e498be04330e91fdf309cca8aa4ff8e49f Mon Sep 17 00:00:00 2001 From: FrankApiyo Date: Thu, 29 Aug 2024 16:09:31 +0300 Subject: [PATCH] Add URLScopedRateThrottle throttle class Co-Authored-By: Kelvin Muchiri --- onadata/libs/tests/test_throttle.py | 23 ++++++++++++++++++++++- onadata/libs/throttle.py | 11 ++++++++++- 2 files changed, 32 insertions(+), 2 deletions(-) diff --git a/onadata/libs/tests/test_throttle.py b/onadata/libs/tests/test_throttle.py index c7b2c2a8e0..97af7bac15 100644 --- a/onadata/libs/tests/test_throttle.py +++ b/onadata/libs/tests/test_throttle.py @@ -1,9 +1,30 @@ from django.core.cache import cache +from django.contrib.auth.models import AnonymousUser from django.test import TestCase, override_settings from rest_framework.test import APIRequestFactory -from onadata.libs.throttle import RequestHeaderThrottle +from onadata.libs.throttle import RequestHeaderThrottle, URLScopedRateThrottle + +class URLScopedThrottlingTests(TestCase): + + + def setUp(self): + """ + Reset the cache so that no throttles will be active + """ + cache.clear() + self.factory = APIRequestFactory() + self.throttle = URLScopedRateThrottle() + + def test_cache_key_is_correctly_set(self): + """ + Cache key is set correctly by anonymous user + """ + request = self.factory.get("/bob/submission") + request.user = AnonymousUser() + key = self.throttle.get_cache_key(request, None) + self.assertEqual(key,'/bob/submission_throttle_None_127.0.0.1') class ThrottlingTests(TestCase): diff --git a/onadata/libs/throttle.py b/onadata/libs/throttle.py index ff79442a0b..5b665cf1b8 100644 --- a/onadata/libs/throttle.py +++ b/onadata/libs/throttle.py @@ -4,7 +4,16 @@ from django.conf import settings -from rest_framework.throttling import SimpleRateThrottle +from rest_framework.throttling import SimpleRateThrottle, ScopedRateThrottle + + +class URLScopedRateThrottle(ScopedRateThrottle): + + def get_cache_key(self, request, view): + """ + Override ScopedRatheThrottle method + """ + return f'{request.path}_{super().get_cache_key(request, view)}' class RequestHeaderThrottle(SimpleRateThrottle):