在 DRF 中测试节流的正确方法是什么？

Question

What is the proper way of testing throttling in DRF? I coulnd't find out any answer to this question on the net. I want to have separate tests for each endpoint since each one has custom requests limits (ScopedRateThrottle).

The important thing is that it can't affect other tests - they have to somehow run without throttling and limiting.

Answer 1

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 .

Answer 2

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')

Answer 3

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)

Answer 4

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'):
...