don't set X-Throttle-Wait-Second header if throttle wait is None

This commit is contained in:
Jeremy Satterfield 2013-08-13 15:31:58 -05:00
parent 999056cde1
commit 1d8a80f5cc
2 changed files with 33 additions and 2 deletions

View File

@ -7,7 +7,7 @@ from django.contrib.auth.models import User
from django.core.cache import cache
from rest_framework.test import APIRequestFactory
from rest_framework.views import APIView
from rest_framework.throttling import UserRateThrottle, ScopedRateThrottle
from rest_framework.throttling import BaseThrottle, UserRateThrottle, ScopedRateThrottle
from rest_framework.response import Response
@ -21,6 +21,14 @@ class User3MinRateThrottle(UserRateThrottle):
scope = 'minutes'
class NonTimeThrottle(BaseThrottle):
def allow_request(self, request, view):
if not hasattr(self.__class__, 'called'):
self.__class__.called = True
return True
return False
class MockView(APIView):
throttle_classes = (User3SecRateThrottle,)
@ -35,6 +43,13 @@ class MockView_MinuteThrottling(APIView):
return Response('foo')
class MockView_NonTimeThrottling(APIView):
throttle_classes = (NonTimeThrottle,)
def get(self, request):
return Response('foo')
class ThrottlingTests(TestCase):
def setUp(self):
"""
@ -140,6 +155,22 @@ class ThrottlingTests(TestCase):
(80, None)
))
def test_non_time_throttle(self):
"""
Ensure for second based throttles.
"""
request = self.factory.get('/')
self.assertFalse(hasattr(MockView_NonTimeThrottling.throttle_classes[0], 'called'))
response = MockView_NonTimeThrottling.as_view()(request)
self.assertFalse('X-Throttle-Wait-Seconds' in response)
self.assertTrue(MockView_NonTimeThrottling.throttle_classes[0].called)
response = MockView_NonTimeThrottling.as_view()(request)
self.assertFalse('X-Throttle-Wait-Seconds' in response)
class ScopedRateThrottleTests(TestCase):
"""

View File

@ -269,7 +269,7 @@ class APIView(View):
Handle any exception that occurs, by returning an appropriate response,
or re-raising the error.
"""
if isinstance(exc, exceptions.Throttled):
if isinstance(exc, exceptions.Throttled) and exc.wait is not None:
# Throttle wait header
self.headers['X-Throttle-Wait-Seconds'] = '%d' % exc.wait