Переопределение настроек для тестов дросселирования Django Rest Framework

Я пытаюсь протестировать пользовательское дросселирование:

def get_user_rate(user):
    # Returns tupple (user plan quota, total seconds in current month)

class SubscriptionDailyRateThrottle(UserRateThrottle):
    # Define a custom scope name to be referenced by DRF in settings.py
    scope = "subscription"

    def __init__(self):
        super().__init__()

    def custom_throttle_success(self):
        """
        Inserts the current request's timestamp along with the key
        into the cache.
        """
        self.history.insert(0, self.now)
        self.cache.set(self.key, self.history, self.duration)
        return True

    def allow_request(self, request, view):
        """
        Override rest_framework.throttling.SimpleRateThrottle.allow_request

        Check to see if the request should be throttled.
        On success calls `throttle_success`.
        On failure calls `throttle_failure`.
        """

        if request.user.is_authenticated:
            limit, duration = get_user_rate(request.user)
            # Override the default from settings.py
            self.duration = duration
            self.num_requests = limit
            
        self.key = self.get_cache_key(request, view)
        if self.key is None:
            return True

        self.history = self.cache.get(self.key, [])
        self.now = self.timer()

        # Drop any requests from the history which have now passed the throttle duration
        while self.history and self.history[-1] <= self.now - self.duration:
            self.history.pop()

        if len(self.history) >= self.num_requests:
            return self.throttle_failure()
        
        return self.custom_throttle_success()

В settings.py я добавил скорость дросселирования по умолчанию 10/секунду просто для безопасности (она передается первой на DEFAULT_THROTTLE_CLASSES):

REST_FRAMEWORK = {
    'DEFAULT_PERMISSION_CLASSES': (
        'rest_framework.permissions.IsAuthenticated',
    ),
    'DEFAULT_RENDERER_CLASSES': (
        'rest_framework.renderers.JSONRenderer',
    ),
    'DEFAULT_AUTHENTICATION_CLASSES': [
        'rest_framework.authentication.TokenAuthentication',
    ],
    'DEFAULT_THROTTLE_CLASSES': [
        'rest_framework.throttling.UserRateThrottle',
        'api.throttling.SubscriptionDailyRateThrottle'
    ],
    'DEFAULT_THROTTLE_RATES': {
        'user': '10/second',
    }
}

Тест, который я хочу написать, очень прост, если у меня есть пользователь с заданным планом, я хочу проверить, что пользователь может сделать до N запросов без дросселирования:

class TestThrottling(TestCase):
    def test_plan_quota(self):
        user = User.objects.create_user(username='test', email='test@email.com', password='test')
        Plan.objects.create(user=user, plan=1) # plan 1 has N requests per month                                                                       
        token, _ = Token.objects.get_or_create(user=user)                
        auth_client = Client(HTTP_AUTHORIZATION='Token ' + token.key)
        
        url = reverse('some_endpoint')

        for k in range(N): # Being N the user plan quota
            response = auth_client.get(url)
            self.assertNotEqual(response.status_code, 429)
        
        response = auth_client.get(url)
        self.assertEqual(response.status_code, 429)

Проблема, которая у меня есть, заключается в скорости по умолчанию 10/сек, которая есть в настройках, потому что она разрывает цикл до достижения квоты пользовательского плана. Я хочу удалить эту скорость по умолчанию из настроек, чтобы проверить, что мое дросселирование работает нормально, я также могу установить таймер, чтобы не делать более 10 запросов в секунду, но квоты планов очень высоки, и это займет несколько часов. Я попробовал переопределить настройки, добавив:

# Override default user throttling
new_config = settings.REST_FRAMEWORK.copy()
new_config['DEFAULT_THROTTLE_CLASSES'] = ['api.throttling.SubscriptionDailyRateThrottle']

@override_settings(REST_FRAMEWORK=new_config)
def test_plan_quota(self):
    ...

Таким образом я мог бы удалить стандартную скорость 10/сек для теста, к сожалению, это не работает, потому что иногда DRF не обновляет настройки. Есть предложения, как это решить?

Проблема в том, что UserRateThrottle всегда будет применяться, поэтому даже если вы заставите тесты работать, подписчики все равно будут привязаны к тарифу по умолчанию UserRateTrottle в производстве.

Решением является наличие обоих типов дросселирования в одном классе:

REST_FRAMEWORK = {
    'DEFAULT_PERMISSION_CLASSES': (
        'rest_framework.permissions.IsAuthenticated',
    ),
    'DEFAULT_RENDERER_CLASSES': (
        'rest_framework.renderers.JSONRenderer',
    ),
    'DEFAULT_AUTHENTICATION_CLASSES': [
        'rest_framework.authentication.TokenAuthentication',
    ],
    'DEFAULT_THROTTLE_CLASSES': [
        'api.throttling.SubscriptionDailyRateThrottle',
    ],
    'DEFAULT_THROTTLE_RATES': {
        'user': '10/second',
    }
}

А затем дифференцировать скорость и продолжительность в зависимости от того, вошел пользователь в систему или нет:

class SubscriptionDailyRateThrottle(UserRateThrottle):
    # Define a custom scope name to be referenced by DRF in settings.py
    scope = "subscription"

    def __init__(self):
        super().__init__()

    def get_user_rate(self, user):
        # Returns tuple (user plan quota, total seconds in current month)
        if user.is_authenticated:
            return (custom_subscriber_rate, custom_subscriber_duration)
        else:
            return (default_rate, default_duration)

    def get_cache_key(self, request, view):
        # Set cache_key scope and ident based on subscriber or anon.
        if request.user and request.user.is_authenticated:
            ident = request.user.pk
            scope = self.scope
        else:
            ident = self.get_ident(request)
            scope = 'anon'

        return self.cache_format % {'scope': scope, 'ident': ident}

    def custom_throttle_success(self):
        """
        Inserts the current request's timestamp along with the key
        into the cache.
        """
        self.history.insert(0, self.now)
        self.cache.set(self.key, self.history, self.duration)
        return True

    def allow_request(self, request, view):
        """
        Override rest_framework.throttling.SimpleRateThrottle.allow_request

        Check to see if the request should be throttled.
        On success calls `throttle_success`.
        On failure calls `throttle_failure`.
        """

        # Get limit and duration for all requests...
        limit, duration = self.get_user_rate(request.user)
        self.duration = duration
        self.num_requests = limit
            
        self.key = self.get_cache_key(request, view)
        if self.key is None:
            return True

        self.history = self.cache.get(self.key, [])
        self.now = self.timer()

        # Drop any requests from the history which have now passed the throttle duration
        while self.history and self.history[-1] <= self.now - self.duration:
            self.history.pop()

        if len(self.history) >= self.num_requests:
            return self.throttle_failure()
        
        return self.custom_throttle_success()
Вернуться на верх