def setUp(self):
"""
Initialises Rate Limit class and delete all keys from Redis.
"""
self.rate_limit = RateLimit(resource='test', client='localhost',
max_requests=10)
self.rate_limit._reset()
def test_context_manager_returns_usage(self):
"""
Should return the usage when used as a context manager.
"""
self.rate_limit = RateLimit(resource='test', client='localhost',
max_requests=1, expire=1)
with self.rate_limit as usage:
self.assertEqual(usage, 1)
def test_wait_time_limit_expired(self):
"""
Should report wait time equal to expire / max_requests before any
requests were made and after the limit has expired.
"""
self.rate_limit = RateLimit(resource='test', client='localhost',
max_requests=10, expire=1)
self.assertEqual(self.rate_limit.get_wait_time(), 1./10)
self._make_10_requests()
time.sleep(1)
self.assertEqual(self.rate_limit.get_wait_time(), 1./10)
def test_acquire_with_blocking(self):
"""
Should not raise TooManyRequests Exception when the quota is over, instead,
must wait for quota to restore and re-acquire from it
"""
self.rate_limit = RateLimit(resource='test',
client='localhost',
max_requests=10)
self._make_10_requests()
with self.rate_limit:
pass
self.assertLessEqual(self.rate_limit.get_usage(), 10)
def test_wait_time_limit_reached(self):
"""
Should report wait time approximately equal to expire after reaching
the limit without delay between requests.
"""
self.rate_limit = RateLimit(resource='test', client='localhost',
max_requests=10, expire=1)
self._make_10_requests()
with self.assertRaises(TooManyRequests):
with self.rate_limit:
pass
self.assertAlmostEqual(self.rate_limit.get_wait_time(), 1, places=2)
def test_not_expired(self):
"""
Should raise TooManyRequests Exception when the expire time has not
been reached yet.
"""
self.rate_limit = RateLimit(resource='test', client='localhost',
max_requests=10, expire=2)
self._make_10_requests()
time.sleep(1)
with self.assertRaises(TooManyRequests):
with self.rate_limit:
pass
def test_ignored_clients(self):
"""
Should not increment counter if client is part of ignored_clients list.
"""
self.rate_limit = RateLimit(resource='test',
client='localhost',
ignored_clients=['localhost'],
max_requests=10,
expire=2)
self.assertEqual(self.rate_limit.get_usage(), 0)
self.assertEqual(self.rate_limit.has_been_reached(), False)
self._make_10_requests()
self.assertEqual(self.rate_limit.get_usage(), 0)
self.assertEqual(self.rate_limit.has_been_reached(), False)
def is_login_rate_limited(account_id):
pool = get_pool()
try:
with RateLimit(resource='login',
client=str(account_id),
max_requests=10,
expire=300,
redis_pool=pool):
return False
except TooManyRequests:
return True
def test_timeout_when_acquire_with_blocking(self):
"""
Should raise Timeout Exception when the quota is over and
we wait too long for it to become available
"""
self.rate_limit = RateLimit(resource='test',
client='localhost',
max_requests=10,
expire=5,
blocking=False)
self.rate_limit_2 = RateLimit(resource='test',
client='localhost',
max_requests=10,
expire=5)
self.rate_limit_2._acquire_check_interval = 3
self.rate_limit_2._acquire_timeout = 5
with self.rate_limit: # start quota
pass
time.sleep(3)
with self.assertRaises(TooManyRequests):
self._make_10_requests() # exceed quota
# By this time there is 100% quota usage and 2 secs before quota recovery
self.assertEqual(self.rate_limit.acquire_attempt, 0)
proc = multiprocessing.Process(target=self.another_thread)
proc.start()
with self.assertRaises(redis_rate_limit.QuotaTimeout):
with self.rate_limit_2:
# Try to acquire. first attempt fails, next attempt in 3 secs
# But after 2 secs quota restores and another instance use it up immediately
# So after another second this instance checks again and see no quota available
# Next time it checks is after another 3 secs 3+3=6 > 5 (max wait time)
pass
proc.join()
def rate_limiter(context):
resource = "%s:%s" % (context.crawler.name, context.stage.name)
rate_limit = context.stage.rate_limit
if settings.REDIS_HOST:
try:
with RateLimit(resource=resource,
client='memorious',
max_requests=rate_limit,
redis_pool=redis_pool):
yield
except TooManyRequests:
raise RateLimitException
else:
if resource in global_call_log:
last_called = global_call_log[resource]
diff = (datetime.datetime.now() - last_called).total_seconds()
if diff < (1.0 / rate_limit):
raise RateLimitException
global_call_log[resource] = datetime.datetime.now()
yield
def test_limit_10_using_rate_limiter(self):
"""
Should raise TooManyRequests Exception when trying to increment for the
eleventh time.
"""
self.rate_limit = RateLimiter(resource='test', max_requests=10,
expire=2).limit(client='localhost')
self.assertEqual(self.rate_limit.get_usage(), 0)
self.assertEqual(self.rate_limit.has_been_reached(), False)
self._make_10_requests()
self.assertEqual(self.rate_limit.get_usage(), 10)
self.assertEqual(self.rate_limit.has_been_reached(), True)
with self.assertRaises(TooManyRequests):
with self.rate_limit:
pass
self.assertEqual(self.rate_limit.get_usage(), 11)
self.assertEqual(self.rate_limit.has_been_reached(), True)
class TestRedisRateLimit(unittest.TestCase):
def setUp(self):
"""
Initialises Rate Limit class and delete all keys from Redis.
"""
self.rate_limit = RateLimit(resource='test',
client='localhost',
max_requests=10)
self.rate_limit._reset()
def _make_10_requests(self):
"""
Increments usage ten times.
"""
for x in range(0, 10):
with self.rate_limit:
pass
def test_limit_10_max_request(self):
"""
Should raise TooManyRequests Exception when trying to increment for the
eleventh time.
"""
self.assertEqual(self.rate_limit.get_usage(), 0)
self.assertEqual(self.rate_limit.has_been_reached(), False)
self._make_10_requests()
self.assertEqual(self.rate_limit.get_usage(), 10)
self.assertEqual(self.rate_limit.has_been_reached(), True)
with self.assertRaises(TooManyRequests):
with self.rate_limit:
pass
self.assertEqual(self.rate_limit.get_usage(), 11)
self.assertEqual(self.rate_limit.has_been_reached(), True)
def test_expire(self):
"""
Should not raise TooManyRequests Exception when trying to increment for
the eleventh time after the expire time.
"""
self._make_10_requests()
time.sleep(1)
with self.rate_limit:
pass
def test_not_expired(self):
"""
Should raise TooManyRequests Exception when the expire time has not
been reached yet.
"""
self.rate_limit = RateLimit(resource='test',
client='localhost',
max_requests=10,
expire=2)
self._make_10_requests()
time.sleep(1)
with self.assertRaises(TooManyRequests):
with self.rate_limit:
pass
def test_limit_10_using_rate_limiter(self):
"""
Should raise TooManyRequests Exception when trying to increment for the
eleventh time.
"""
self.rate_limit = RateLimiter(resource='test',
max_requests=10,
expire=2).limit(client='localhost')
self.assertEqual(self.rate_limit.get_usage(), 0)
self.assertEqual(self.rate_limit.has_been_reached(), False)
self._make_10_requests()
self.assertEqual(self.rate_limit.get_usage(), 10)
self.assertEqual(self.rate_limit.has_been_reached(), True)
with self.assertRaises(TooManyRequests):
with self.rate_limit:
pass
self.assertEqual(self.rate_limit.get_usage(), 11)
self.assertEqual(self.rate_limit.has_been_reached(), True)
def test_wait_time_limit_reached(self):
"""
Should report wait time approximately equal to expire after reaching
the limit without delay between requests.
"""
self.rate_limit = RateLimit(resource='test',
client='localhost',
max_requests=10,
expire=1)
self._make_10_requests()
with self.assertRaises(TooManyRequests):
with self.rate_limit:
pass
self.assertAlmostEqual(self.rate_limit.get_wait_time(), 1, places=2)
def test_wait_time_limit_expired(self):
"""
Should report wait time equal to expire / max_requests before any
requests were made and after the limit has expired.
"""
self.rate_limit = RateLimit(resource='test',
client='localhost',
max_requests=10,
expire=1)
self.assertEqual(self.rate_limit.get_wait_time(), 1. / 10)
self._make_10_requests()
time.sleep(1)
self.assertEqual(self.rate_limit.get_wait_time(), 1. / 10)
class TestRedisRateLimit(unittest.TestCase):
def setUp(self):
"""
Initialises Rate Limit class and delete all keys from Redis.
"""
self.rate_limit = RateLimit(resource='test',
client='localhost',
blocking=False,
max_requests=10)
self.rate_limit._reset()
def _make_10_requests(self):
"""
Increments usage ten times.
"""
for x in range(0, 10):
with self.rate_limit:
pass
def test_limit_10_max_request(self):
"""
Should raise TooManyRequests Exception when trying to increment for the
eleventh time.
"""
self.assertEqual(self.rate_limit.get_usage(), 0)
self.assertEqual(self.rate_limit.has_been_reached(), False)
self._make_10_requests()
self.assertEqual(self.rate_limit.get_usage(), 10)
self.assertEqual(self.rate_limit.has_been_reached(), True)
with self.assertRaises(TooManyRequests):
with self.rate_limit:
pass
self.assertEqual(self.rate_limit.get_usage(), 10)
self.assertEqual(self.rate_limit.has_been_reached(), True)
def test_expire(self):
"""
Should not raise TooManyRequests Exception when trying to increment for
the eleventh time after the expire time.
"""
self._make_10_requests()
time.sleep(1)
with self.rate_limit:
pass
def test_not_expired(self):
"""
Should raise TooManyRequests Exception when the expire time has not
been reached yet.
"""
self.rate_limit = RateLimit(resource='test',
client='localhost',
blocking=False,
max_requests=10,
expire=2)
self._make_10_requests()
time.sleep(1)
with self.assertRaises(TooManyRequests):
with self.rate_limit:
pass
def test_acquire_with_blocking(self):
"""
Should not raise TooManyRequests Exception when the quota is over, instead,
must wait for quota to restore and re-acquire from it
"""
self.rate_limit = RateLimit(resource='test',
client='localhost',
max_requests=10)
self._make_10_requests()
with self.rate_limit:
pass
self.assertLessEqual(self.rate_limit.get_usage(), 10)
def another_thread(self):
self.rate_limit.blocking = True
self._make_10_requests()
def test_timeout_when_acquire_with_blocking(self):
"""
Should raise Timeout Exception when the quota is over and
we wait too long for it to become available
"""
self.rate_limit = RateLimit(resource='test',
client='localhost',
max_requests=10,
expire=5,
blocking=False)
self.rate_limit_2 = RateLimit(resource='test',
client='localhost',
max_requests=10,
expire=5)
self.rate_limit_2._acquire_check_interval = 3
self.rate_limit_2._acquire_timeout = 5
with self.rate_limit: # start quota
pass
time.sleep(3)
with self.assertRaises(TooManyRequests):
self._make_10_requests() # exceed quota
# By this time there is 100% quota usage and 2 secs before quota recovery
self.assertEqual(self.rate_limit.acquire_attempt, 0)
proc = multiprocessing.Process(target=self.another_thread)
proc.start()
with self.assertRaises(redis_rate_limit.QuotaTimeout):
with self.rate_limit_2:
# Try to acquire. first attempt fails, next attempt in 3 secs
# But after 2 secs quota restores and another instance use it up immediately
# So after another second this instance checks again and see no quota available
# Next time it checks is after another 3 secs 3+3=6 > 5 (max wait time)
pass
proc.join()
def test_nested_not_allowed(self):
"""
Should raise GaveUp Exception when tryig to enter the context when already in the context for the same instance
"""
with self.assertRaises(redis_rate_limit.GaveUp):
with self.rate_limit:
with self.rate_limit:
pass
def test_acquire_callbacks(self):
def cb(rl):
rl.callback_called = True
self.rate_limit.add_post_enter_callback(cb)
with self.rate_limit:
self.assertEquals(self.rate_limit.callback_called, True)
self.assertEquals(self.rate_limit.post_enter_callbacks, set())
class TestRedisRateLimit(unittest.TestCase):
def setUp(self):
"""
Initialises Rate Limit class and delete all keys from Redis.
"""
self.rate_limit = RateLimit(resource='test',
client='localhost',
max_requests=10)
self.rate_limit._reset()
def _make_10_requests(self):
"""
Increments usage ten times.
"""
for x in range(0, 10):
with self.rate_limit:
pass
def test_limit_10_max_request(self):
"""
Should raise TooManyRequests Exception when trying to increment for the
eleventh time.
"""
self.assertEqual(self.rate_limit.get_usage(), 0)
self.assertEqual(self.rate_limit.has_been_reached(), False)
self._make_10_requests()
self.assertEqual(self.rate_limit.get_usage(), 10)
self.assertEqual(self.rate_limit.has_been_reached(), True)
with self.assertRaises(TooManyRequests):
with self.rate_limit:
pass
self.assertEqual(self.rate_limit.get_usage(), 11)
self.assertEqual(self.rate_limit.has_been_reached(), True)
def test_expire(self):
"""
Should not raise TooManyRequests Exception when trying to increment for
the eleventh time after the expire time.
"""
self._make_10_requests()
time.sleep(1)
with self.rate_limit:
pass
def test_not_expired(self):
"""
Should raise TooManyRequests Exception when the expire time has not
been reached yet.
"""
self.rate_limit = RateLimit(resource='test',
client='localhost',
max_requests=10,
expire=2)
self._make_10_requests()
time.sleep(1)
with self.assertRaises(TooManyRequests):
with self.rate_limit:
pass
def test_limit_10_using_rate_limiter(self):
"""
Should raise TooManyRequests Exception when trying to increment for the
eleventh time.
"""
self.rate_limit = RateLimiter(resource='test',
max_requests=10,
expire=2).limit(client='localhost')
self.assertEqual(self.rate_limit.get_usage(), 0)
self.assertEqual(self.rate_limit.has_been_reached(), False)
self._make_10_requests()
self.assertEqual(self.rate_limit.get_usage(), 10)
self.assertEqual(self.rate_limit.has_been_reached(), True)
with self.assertRaises(TooManyRequests):
with self.rate_limit:
pass
self.assertEqual(self.rate_limit.get_usage(), 11)
self.assertEqual(self.rate_limit.has_been_reached(), True)
def test_wait_time_limit_reached(self):
"""
Should report wait time approximately equal to expire after reaching
the limit without delay between requests.
"""
self.rate_limit = RateLimit(resource='test',
client='localhost',
max_requests=10,
expire=1)
self._make_10_requests()
with self.assertRaises(TooManyRequests):
with self.rate_limit:
pass
self.assertAlmostEqual(self.rate_limit.get_wait_time(), 1, places=1)
def test_wait_time_limit_expired(self):
"""
Should report wait time equal to expire / max_requests before any
requests were made and after the limit has expired.
"""
self.rate_limit = RateLimit(resource='test',
client='localhost',
max_requests=10,
expire=1)
self.assertEqual(self.rate_limit.get_wait_time(), 1. / 10)
self._make_10_requests()
time.sleep(1)
self.assertEqual(self.rate_limit.get_wait_time(), 1. / 10)
def test_context_manager_returns_usage(self):
"""
Should return the usage when used as a context manager.
"""
self.rate_limit = RateLimit(resource='test',
client='localhost',
max_requests=1,
expire=1)
with self.rate_limit as usage:
self.assertEqual(usage, 1)
def test_decorator(self):
@RateLimit(resource='test',
client='localhost',
max_requests=1,
expire=1)
def foo():
pass
foo()
with self.assertRaises(TooManyRequests):
foo()
def test_sleep_rate_limit(self):
sleep_time = 0.2
@SleepRateLimit(resource='test',
client='localhost',
max_requests=1,
expire=10,
sleep_time=sleep_time)
def foo():
pass
t0 = time.time()
foo()
t1 = time.time()
self.assertAlmostEqual(0, t1 - t0, places=1)
foo()
t2 = time.time()
self.assertAlmostEqual(sleep_time, t2 - t1, places=2)
foo()
t3 = time.time()
self.assertAlmostEqual(sleep_time, t3 - t2, places=2)
foo()
t4 = time.time()
self.assertAlmostEqual(2 * sleep_time, t4 - t3, places=2)
class TestRedisRateLimit(unittest.TestCase):
def setUp(self):
"""
Initialises Rate Limit class and delete all keys from Redis.
"""
self.rate_limit = RateLimit(resource='test',
client='localhost',
max_requests=10)
self.rate_limit._reset()
def _make_10_requests(self):
"""
Increments usage ten times.
"""
for x in range(0, 10):
with self.rate_limit:
pass
def test_limit_10_max_request(self):
"""
Should raise TooManyRequests Exception when trying to increment for the
eleventh time.
"""
self.assertEqual(self.rate_limit.get_usage(), 0)
self.assertEqual(self.rate_limit.has_been_reached(), False)
self._make_10_requests()
self.assertEqual(self.rate_limit.get_usage(), 10)
self.assertEqual(self.rate_limit.has_been_reached(), True)
with self.assertRaises(TooManyRequests):
with self.rate_limit:
pass
self.assertEqual(self.rate_limit.get_usage(), 11)
self.assertEqual(self.rate_limit.has_been_reached(), True)
def test_expire(self):
"""
Should not raise TooManyRequests Exception when trying to increment for
the eleventh time after the expire time.
"""
self._make_10_requests()
time.sleep(1)
with self.rate_limit:
pass
def test_not_expired(self):
"""
Should raise TooManyRequests Exception when the expire time has not
been reached yet.
"""
self.rate_limit = RateLimit(resource='test',
client='localhost',
max_requests=10,
expire=2)
self._make_10_requests()
time.sleep(1)
with self.assertRaises(TooManyRequests):
with self.rate_limit:
pass
class TestRedisRateLimit(unittest.TestCase):
def setUp(self):
"""
Initialises Rate Limit class and delete all keys from Redis.
"""
self.rate_limit = RateLimit(resource='test',
client='localhost',
max_requests=10)
self.rate_limit._reset()
def _make_10_requests(self):
"""
Increments usage ten times.
"""
for x in range(0, 10):
with self.rate_limit:
pass
def test_limit_10_max_request(self):
"""
Should raise TooManyRequests Exception when trying to increment for the
eleventh time.
"""
self.assertEqual(self.rate_limit.get_usage(), 0)
self.assertEqual(self.rate_limit.has_been_reached(), False)
self._make_10_requests()
self.assertEqual(self.rate_limit.get_usage(), 10)
self.assertEqual(self.rate_limit.has_been_reached(), True)
with self.assertRaises(TooManyRequests):
with self.rate_limit:
pass
self.assertEqual(self.rate_limit.get_usage(), 11)
self.assertEqual(self.rate_limit.has_been_reached(), True)
def test_ignored_clients(self):
"""
Should not increment counter if client is part of ignored_clients list.
"""
self.rate_limit = RateLimit(resource='test',
client='localhost',
ignored_clients=['localhost'],
max_requests=10,
expire=2)
self.assertEqual(self.rate_limit.get_usage(), 0)
self.assertEqual(self.rate_limit.has_been_reached(), False)
self._make_10_requests()
self.assertEqual(self.rate_limit.get_usage(), 0)
self.assertEqual(self.rate_limit.has_been_reached(), False)
def test_expire(self):
"""
Should not raise TooManyRequests Exception when trying to increment for
the eleventh time after the expire time.
"""
self._make_10_requests()
time.sleep(1)
with self.rate_limit:
pass
def test_not_expired(self):
"""
Should raise TooManyRequests Exception when the expire time has not
been reached yet.
"""
self.rate_limit = RateLimit(resource='test',
client='localhost',
max_requests=10,
expire=2)
self._make_10_requests()
time.sleep(1)
with self.assertRaises(TooManyRequests):
with self.rate_limit:
pass
def test_limit_10_using_rate_limiter(self):
"""
Should raise TooManyRequests Exception when trying to increment for the
eleventh time.
"""
self.rate_limit = RateLimiter(resource='test',
max_requests=10,
expire=2).limit(client='localhost')
self.assertEqual(self.rate_limit.get_usage(), 0)
self.assertEqual(self.rate_limit.has_been_reached(), False)
self._make_10_requests()
self.assertEqual(self.rate_limit.get_usage(), 10)
self.assertEqual(self.rate_limit.has_been_reached(), True)
with self.assertRaises(TooManyRequests):
with self.rate_limit:
pass
self.assertEqual(self.rate_limit.get_usage(), 11)
self.assertEqual(self.rate_limit.has_been_reached(), True)
def test_wait_time_limit_reached(self):
"""
Should report wait time approximately equal to expire after reaching
the limit without delay between requests.
"""
self.rate_limit = RateLimit(resource='test',
client='localhost',
max_requests=10,
expire=1)
self._make_10_requests()
with self.assertRaises(TooManyRequests):
with self.rate_limit:
pass
self.assertAlmostEqual(self.rate_limit.get_wait_time(), 1, places=2)
def test_wait_time_limit_expired(self):
"""
Should report wait time equal to expire / max_requests before any
requests were made and after the limit has expired.
"""
self.rate_limit = RateLimit(resource='test',
client='localhost',
max_requests=10,
expire=1)
self.assertEqual(self.rate_limit.get_wait_time(), 1. / 10)
self._make_10_requests()
time.sleep(1)
self.assertEqual(self.rate_limit.get_wait_time(), 1. / 10)
def test_context_manager_returns_usage(self):
"""
Should return the usage when used as a context manager.
"""
self.rate_limit = RateLimit(resource='test',
client='localhost',
max_requests=1,
expire=1)
with self.rate_limit as usage:
self.assertEqual(usage, 1)
def test_limit_10_using_as_decorator(self):
"""
Should raise TooManyRequests Exception when trying to increment for the
eleventh time.
"""
self.assertEqual(self.rate_limit.get_usage(), 0)
self.assertEqual(self.rate_limit.has_been_reached(), False)
self._make_10_requests()
self.assertEqual(self.rate_limit.get_usage(), 10)
self.assertEqual(self.rate_limit.has_been_reached(), True)
@self.rate_limit
def limit_with_decorator():
pass
with self.assertRaises(TooManyRequests):
limit_with_decorator()
self.assertEqual(self.rate_limit.get_usage(), 11)
self.assertEqual(self.rate_limit.has_been_reached(), True)
def test_increment_multiple(self):
"""
Test incrementing usage by a value > 1
"""
self.rate_limit.increment_usage(7)
self.rate_limit.increment_usage(3)
self.assertEqual(self.rate_limit.get_usage(), 10)
self.assertEqual(self.rate_limit.has_been_reached(), True)
with self.assertRaises(TooManyRequests):
self.rate_limit.increment_usage(1)
self.assertEqual(self.rate_limit.get_usage(), 11)
self.assertEqual(self.rate_limit.has_been_reached(), True)
with self.assertRaises(TooManyRequests):
self.rate_limit.increment_usage(5)
self.assertEqual(self.rate_limit.get_usage(), 16)
self.assertEqual(self.rate_limit.has_been_reached(), True)
def test_increment_multiple_too_much(self):
"""
Test that we cannot bulk-increment a value higher than
the bucket limit.
"""
with self.assertRaises(ValueError):
self.rate_limit.increment_usage(11)
self.assertEqual(self.rate_limit.get_usage(), 0)
self.assertEqual(self.rate_limit.has_been_reached(), False)
def test_increment_by_zero(self):
"""
Should not allow increment by zero.
"""
self.assertEqual(self.rate_limit.get_usage(), 0)
self.assertEqual(self.rate_limit.has_been_reached(), False)
self.rate_limit.increment_usage(5)
self.assertEqual(self.rate_limit.get_usage(), 5)
self.assertEqual(self.rate_limit.has_been_reached(), False)
with self.assertRaises(ValueError):
self.rate_limit.increment_usage(0)
self.assertEqual(self.rate_limit.get_usage(), 5)
self.assertEqual(self.rate_limit.has_been_reached(), False)
def test_increment_by_negative(self):
"""
Should not allow decrement the counter.
"""
self.assertEqual(self.rate_limit.get_usage(), 0)
self.assertEqual(self.rate_limit.has_been_reached(), False)
with self.assertRaises(ValueError):
self.rate_limit.increment_usage(-5)
self.assertEqual(self.rate_limit.get_usage(), 0)
self.assertEqual(self.rate_limit.has_been_reached(), False)
