Changeset View
Changeset View
Standalone View
Standalone View
swh/web/tests/common/test_throttling.py
| Show All 12 Lines | |||||
| from django.core.cache import cache | from django.core.cache import cache | ||||
| from django.test.utils import override_settings | from django.test.utils import override_settings | ||||
| from rest_framework.views import APIView | from rest_framework.views import APIView | ||||
| from rest_framework.response import Response | from rest_framework.response import Response | ||||
| from rest_framework.test import APIRequestFactory | from rest_framework.test import APIRequestFactory | ||||
| from rest_framework.decorators import api_view | from rest_framework.decorators import api_view | ||||
| from nose.tools import istest, nottest | |||||
| from swh.web.common.throttling import ( | from swh.web.common.throttling import ( | ||||
| SwhWebRateThrottle, throttle_scope | SwhWebRateThrottle, throttle_scope | ||||
| ) | ) | ||||
| from swh.web.tests.testcase import SWHWebTestCase | from swh.web.tests.testcase import SWHWebTestCase | ||||
| class MockViewScope1(APIView): | class MockViewScope1(APIView): | ||||
| ▲ Show 20 Lines • Show All 48 Lines • ▼ Show 20 Lines | |||||
| class ThrottlingTests(SWHWebTestCase): | class ThrottlingTests(SWHWebTestCase): | ||||
| def setUp(self): | def setUp(self): | ||||
| """ | """ | ||||
| Reset the cache so that no throttles will be active | Reset the cache so that no throttles will be active | ||||
| """ | """ | ||||
| cache.clear() | cache.clear() | ||||
| self.factory = APIRequestFactory() | self.factory = APIRequestFactory() | ||||
| @nottest | |||||
| def check_response(self, response, status_code, | def check_response(self, response, status_code, | ||||
| limit=None, remaining=None): | limit=None, remaining=None): | ||||
| assert response.status_code == status_code | assert response.status_code == status_code | ||||
| if limit is not None: | if limit is not None: | ||||
| assert response['X-RateLimit-Limit'] == str(limit) | assert response['X-RateLimit-Limit'] == str(limit) | ||||
| else: | else: | ||||
| assert 'X-RateLimit-Limit' not in response | assert 'X-RateLimit-Limit' not in response | ||||
| if remaining is not None: | if remaining is not None: | ||||
| assert response['X-RateLimit-Remaining'] == str(remaining) | assert response['X-RateLimit-Remaining'] == str(remaining) | ||||
| else: | else: | ||||
| assert 'X-RateLimit-Remaining' not in response | assert 'X-RateLimit-Remaining' not in response | ||||
| @istest | def test_scope1_requests_are_throttled(self): | ||||
| def scope1_requests_are_throttled(self): | |||||
| """ | """ | ||||
| Ensure request rate is limited in scope1 | Ensure request rate is limited in scope1 | ||||
| """ | """ | ||||
| for i in range(scope1_limiter_rate): | for i in range(scope1_limiter_rate): | ||||
| response = self.client.get('/scope1_class') | response = self.client.get('/scope1_class') | ||||
| self.check_response(response, 200, scope1_limiter_rate, | self.check_response(response, 200, scope1_limiter_rate, | ||||
| scope1_limiter_rate - i - 1) | scope1_limiter_rate - i - 1) | ||||
| response = self.client.get('/scope1_class') | response = self.client.get('/scope1_class') | ||||
| self.check_response(response, 429, scope1_limiter_rate, 0) | self.check_response(response, 429, scope1_limiter_rate, 0) | ||||
| for i in range(scope1_limiter_rate_post): | for i in range(scope1_limiter_rate_post): | ||||
| response = self.client.post('/scope1_class') | response = self.client.post('/scope1_class') | ||||
| self.check_response(response, 200, scope1_limiter_rate_post, | self.check_response(response, 200, scope1_limiter_rate_post, | ||||
| scope1_limiter_rate_post - i - 1) | scope1_limiter_rate_post - i - 1) | ||||
| response = self.client.post('/scope1_class') | response = self.client.post('/scope1_class') | ||||
| self.check_response(response, 429, scope1_limiter_rate_post, 0) | self.check_response(response, 429, scope1_limiter_rate_post, 0) | ||||
| @istest | def test_scope2_requests_are_throttled(self): | ||||
| def scope2_requests_are_throttled(self): | |||||
| """ | """ | ||||
| Ensure request rate is limited in scope2 | Ensure request rate is limited in scope2 | ||||
| """ | """ | ||||
| for i in range(scope2_limiter_rate): | for i in range(scope2_limiter_rate): | ||||
| response = self.client.get('/scope2_func') | response = self.client.get('/scope2_func') | ||||
| self.check_response(response, 200, scope2_limiter_rate, | self.check_response(response, 200, scope2_limiter_rate, | ||||
| scope2_limiter_rate - i - 1) | scope2_limiter_rate - i - 1) | ||||
| response = self.client.get('/scope2_func') | response = self.client.get('/scope2_func') | ||||
| self.check_response(response, 429, scope2_limiter_rate, 0) | self.check_response(response, 429, scope2_limiter_rate, 0) | ||||
| for i in range(scope2_limiter_rate_post): | for i in range(scope2_limiter_rate_post): | ||||
| response = self.client.post('/scope2_func') | response = self.client.post('/scope2_func') | ||||
| self.check_response(response, 200, scope2_limiter_rate_post, | self.check_response(response, 200, scope2_limiter_rate_post, | ||||
| scope2_limiter_rate_post - i - 1) | scope2_limiter_rate_post - i - 1) | ||||
| response = self.client.post('/scope2_func') | response = self.client.post('/scope2_func') | ||||
| self.check_response(response, 429, scope2_limiter_rate_post, 0) | self.check_response(response, 429, scope2_limiter_rate_post, 0) | ||||
| @istest | def test_scope3_requests_are_throttled_exempted(self): | ||||
| def scope3_requests_are_throttled_exempted(self): | |||||
| """ | """ | ||||
| Ensure request rate is not limited in scope3 as | Ensure request rate is not limited in scope3 as | ||||
| requests coming from localhost are exempted from rate limit. | requests coming from localhost are exempted from rate limit. | ||||
| """ | """ | ||||
| for _ in range(scope3_limiter_rate+1): | for _ in range(scope3_limiter_rate+1): | ||||
| response = self.client.get('/scope3_class') | response = self.client.get('/scope3_class') | ||||
| self.check_response(response, 200) | self.check_response(response, 200) | ||||
| Show All 11 Lines | |||||