[Fixed]-What is the proper way of testing throttling in DRF?

by

6👍

✅

Like people already mentioned, this doesn’t exactly fall within the scope of unit tests, but still, how about simply doing something like this:

from django.core.urlresolvers import reverse
from django.test import override_settings
from rest_framework.test import APITestCase, APIClient


class ThrottleApiTests(APITestCase):
    # make sure to override your settings for testing
    TESTING_THRESHOLD = '5/min'
    # THROTTLE_THRESHOLD is the variable that you set for DRF DEFAULT_THROTTLE_RATES
    @override_settings(THROTTLE_THRESHOLD=TESTING_THRESHOLD)
    def test_check_health(self):
        client = APIClient()
        # some end point you want to test (in this case it's a public enpoint that doesn't require authentication
        _url = reverse('check-health')
        # this is probably set in settings in you case
        for i in range(0, self.TESTING_THRESHOLD):
            client.get(_url)

        # this call should err
        response = client.get(_url)
        # 429 - too many requests
        self.assertEqual(response.status_code, 429)

Also, regarding your concerns of side-effects, as long as you do user creation in setUp or setUpTestData, tests will be isolated (as they should), so no need to worry about ‘dirty’ data or scope in that sense.

Regarding cache clearing between tests, I would just add cache.clear() in tearDown or try and clear the specific key defined for throttling.

👤Daniel Dror

7👍

An easy solution is to patch the get_rate method of your throttle class. Thanks to tprestegard for this comment!

I have a custom class in my case:

from rest_framework.throttling import UserRateThrottle

class AuthRateThrottle(UserRateThrottle):
    scope = 'auth'

In your tests:

from unittest.mock import patch
from django.core.cache import cache
from rest_framework import status

class Tests(SimpleTestCase):
    def setUp(self):
        cache.clear()

    @patch('path.to.AuthRateThrottle.get_rate')
    def test_throttling(self, mock):
        mock.return_value = '1/day'
        response = self.client.post(self.url, {})
        self.assertEqual(
            response.status_code,
            status.HTTP_400_BAD_REQUEST,  # some fields are required
        )
        response = self.client.post(self.url, {})
        self.assertEqual(
            response.status_code,
            status.HTTP_429_TOO_MANY_REQUESTS,
        )

It is also possible to patch the method in the DRF package to change the behavior of the standard throttle classes: @patch('rest_framework.throttling.SimpleRateThrottle.get_rate')

👤yofee

0👍

I implemented my own caching mechanism for throttling based on the user and the parameters with which a request is called. You can override SimpleRateThrottle.get_cache_key to get this behavior.

Take this throttle class for example:

class YourCustomThrottleClass(SimpleRateThrottle):
    rate = "1/d"
    scope = "your-custom-throttle"

    def get_cache_key(self, request: Request, view: viewsets.ModelViewSet):
        # we want to throttle the based on the request user as well as the parameter
        # `foo` (i.e. the user can make a request with a different `foo` as many times
        # as they want in a day, but only once a day for a given `foo`).
        foo_request_param = view.kwargs["foo"]
        ident = f"{request.user.pk}_{foo_request_param}"
        # below format is copied from `UserRateThrottle.get_cache_key`.
        return self.cache_format % {"scope": self.scope, "ident": ident}

In order to clear this in a TestCase I call the following method in each test method as required:

def _clear_throttle_cache(self, request_user, foo_param):
    # we need to clear the cache of the throttle limits already stored there.
    throttle = YourCustomThrottleClass()

    # in the below two lines mock whatever attributes on the request and
    # view instances are used to calculate the cache key in `.get_cache_key`
    # which you overrode. Here we use `request.user` and `view.kwargs["foo"]` 
    # to calculate the throttle key, so we mock those.
    pretend_view = MagicMock(kwargs={foo: foo_param})
    pretend_request = MagicMock(user=request_user)
    
    # this is the method you overrode in `YourCustomThrottleClass`.
    throttle_key = throttle.get_cache_key(pretend_request, pretend_view)
    throttle.cache.delete(user_key)
👤Adam Strike

0👍

This is an amendment on yofee’s post which got me 90% there. When using a throttle, custom or otherwise, with a set rate, get_rate is never called. As shown below from the source.

    def __init__(self):
        if not getattr(self, 'rate', None):
            self.rate = self.get_rate()

Hence when one is mocking a throttle with a set rate that is not None, I would recommend patching the rate attribute directly.

...
with mock.patch.object(AuthRateThrottle, 'rate', '1/day'):
...
👤af3ld

Leave a comment