Переопределение настроек для тестов дросселирования 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()