Slightly clean throttling tests

This commit is contained in:
Aarni Koskela 2019-09-26 18:07:32 +03:00
parent 5b990d4092
commit 0bcb275281

View File

@ -20,6 +20,16 @@ from rest_framework.throttling import (
from rest_framework.views import APIView from rest_framework.views import APIView
class ThrottleTestTimerMixin:
"""
Mixin to make throttle timers programmable with `TIMER_SECONDS` for tests.
"""
TIMER_SECONDS = 0
def timer(self):
return self.TIMER_SECONDS
class User3SecRateThrottle(UserRateThrottle): class User3SecRateThrottle(UserRateThrottle):
rate = '3/sec' rate = '3/sec'
scope = 'seconds' scope = 'seconds'
@ -192,7 +202,7 @@ class ThrottlingTests(TestCase):
if expect is not None: if expect is not None:
assert response['Retry-After'] == expect assert response['Retry-After'] == expect
else: else:
assert not'Retry-After' in response assert 'Retry-After' not in response
def test_seconds_fields(self): def test_seconds_fields(self):
""" """
@ -236,9 +246,6 @@ class ThrottlingTests(TestCase):
) )
def test_non_time_throttle(self): def test_non_time_throttle(self):
"""
Ensure for second based throttles.
"""
request = self.factory.get('/') request = self.factory.get('/')
self.assertFalse(hasattr(MockView_NonTimeThrottling.throttle_classes[0], 'called')) self.assertFalse(hasattr(MockView_NonTimeThrottling.throttle_classes[0], 'called'))
@ -260,13 +267,9 @@ class ScopedRateThrottleTests(TestCase):
def setUp(self): def setUp(self):
self.throttle = ScopedRateThrottle() self.throttle = ScopedRateThrottle()
class XYScopedRateThrottle(ScopedRateThrottle): class XYScopedRateThrottle(ThrottleTestTimerMixin, ScopedRateThrottle):
TIMER_SECONDS = 0
THROTTLE_RATES = {'x': '3/min', 'y': '1/min'} THROTTLE_RATES = {'x': '3/min', 'y': '1/min'}
def timer(self):
return self.TIMER_SECONDS
class XView(APIView): class XView(APIView):
throttle_classes = (XYScopedRateThrottle,) throttle_classes = (XYScopedRateThrottle,)
throttle_scope = 'x' throttle_scope = 'x'
@ -375,12 +378,8 @@ class ScopedRateThrottleTests(TestCase):
class XffTestingBase(TestCase): class XffTestingBase(TestCase):
def setUp(self): def setUp(self):
class Throttle(ScopedRateThrottle): class Throttle(ThrottleTestTimerMixin, ScopedRateThrottle):
THROTTLE_RATES = {'test_limit': '1/day'} THROTTLE_RATES = {'test_limit': '1/day'}
TIMER_SECONDS = 0
def timer(self):
return self.TIMER_SECONDS
class View(APIView): class View(APIView):
throttle_classes = (Throttle,) throttle_classes = (Throttle,)