desec-stack/api/desecapi/tests/test_throttling.py

from unittest import mock
import time

from django.core.cache import cache
from django.test import TestCase, override_settings
from rest_framework import status
from rest_framework.response import Response
from rest_framework.views import APIView
from rest_framework.test import APIRequestFactory


def override_rates(rates):
    return override_settings(REST_FRAMEWORK={'DEFAULT_THROTTLE_CLASSES': ['desecapi.throttling.ScopedRatesThrottle'],
                                             'DEFAULT_THROTTLE_RATES': {'test_scope': rates}})


class MockView(APIView):
    throttle_scope = 'test_scope'

    @property
    def throttle_classes(self):
        # Need to import here so that the module is only loaded once the settings override is in effect
        from desecapi.throttling import ScopedRatesThrottle
        return (ScopedRatesThrottle,)

    def get(self, request):
        return Response('foo')


class ThrottlingTestCase(TestCase):
    """
    Based on DRF's test_throttling.py.
    """
    def setUp(self):
        super().setUp()
        self.factory = APIRequestFactory()

    def _test_requests_are_throttled(self, rates, counts, buckets=None):
        def do_test():
            view = MockView.as_view()
            sum_delay = 0
            for delay, count, max_wait in counts:
                sum_delay += delay
                with mock.patch('desecapi.throttling.ScopedRatesThrottle.timer', return_value=time.time() + sum_delay):
                    for _ in range(count):
                        response = view(request)
                        self.assertEqual(response.status_code, status.HTTP_200_OK)

                    response = view(request)
                    self.assertEqual(response.status_code, status.HTTP_429_TOO_MANY_REQUESTS)
                    self.assertTrue(max_wait - 1 <= float(response['Retry-After']) <= max_wait)

        cache.clear()
        request = self.factory.get('/')
        with override_rates(rates):
            do_test()
            if buckets is not None:
                for bucket in buckets:
                    MockView.throttle_scope_bucket = bucket
                    do_test()

    def test_requests_are_throttled_4sec(self):
        self._test_requests_are_throttled(['4/sec'], [(0, 4, 1), (1, 4, 1)])

    def test_requests_are_throttled_4min(self):
        self._test_requests_are_throttled(['4/min'], [(0, 4, 60)])

    def test_requests_are_throttled_multiple(self):
        self._test_requests_are_throttled(['5/s', '4/day'], [(0, 4, 86400)])
        self._test_requests_are_throttled(['4/s', '5/day'], [(0, 4, 1)])

    def test_requests_are_throttled_multiple_cascade(self):
        # We test that we can do 4 requests in the first second and only 2 in the second second
        self._test_requests_are_throttled(['4/s', '6/day'], [(0, 4, 1), (1, 2, 86400)])

    def test_requests_are_throttled_multiple_cascade_with_buckets(self):
        # We test that we can do 4 requests in the first second and only 2 in the second second
        self._test_requests_are_throttled(['4/s', '6/day'], [(0, 4, 1), (1, 2, 86400)], buckets=['foo', 'bar'])