def test_init_arguments(self):
# Test some cases that should raise
self.assertRaises(TypeError, rate_limiter.RateLimiter, 10, "dog")
self.assertRaises(TypeError, rate_limiter.RateLimiter, "cat", 1)
self.assertRaises(ValueError, rate_limiter.RateLimiter, 10, 0)
self.assertRaises(ValueError, rate_limiter.RateLimiter, 0, 5)
self.assertRaises(ValueError, rate_limiter.RateLimiter, 10, -5)
self.assertRaises(ValueError, rate_limiter.RateLimiter, -5, 5)
# And some that should be OK
self.assertIsInstance(rate_limiter.RateLimiter(1, 2),
rate_limiter.RateLimiter)
self.assertIsInstance(rate_limiter.RateLimiter(3, 4),
rate_limiter.RateLimiter)
self.assertIsInstance(rate_limiter.RateLimiter(10000, 1000000),
rate_limiter.RateLimiter)
def test_time_went_backwards(self, new_limit_period_mock, now_mock):
"""Make certain that if the system time went backwards for any reason,
RateLimiter will handle it. Otherwise RateLimiter would break by
not starting a new limiting period for possibly a very long time.
"""
now_mock.return_value = START_DATETIME
rl_obj = rate_limiter.RateLimiter(3, 5)
self.assertEqual(START_DATETIME, rl_obj.current_limit_period_start)
now_mock.return_value = PAST_DATETIME
rl_obj.withdraw_items()
self.assertTrue(new_limit_period_mock.called)
def get_configured_rate_limiter(name, rate_limit, limiting_period):
if rate_limit == 0:
LOG.info("%s directives will not be rate limited during this run." %
(name))
return None
LOG.info("%(name)s directives will be rate limited to %(rate_limit)d "
"every %(seconds)d second(s).", {
'name': name,
'rate_limit': rate_limit,
'seconds': limiting_period})
return rate_limiter.RateLimiter(
limit=rate_limit,
limit_period=limiting_period)
def test_update_limit_period(self, now_mock):
now_mock.return_value = START_DATETIME
rl_obj = rate_limiter.RateLimiter(10, 5)
self.assertEqual(START_DATETIME, rl_obj.current_limit_period_start)
now_mock.return_value = START_DATETIME + ONE_SECOND
rl_obj.current_count = 3
rl_obj._update_limit_period()
# These shouldn't change yet.
self.assertEqual(START_DATETIME, rl_obj.current_limit_period_start)
self.assertEqual(3, rl_obj.current_count)
# Now we're beyond the initial period, and the rate limiter should
# create a new period starting when withdraw_items() is called, and
# reset the count.
now_mock.return_value = START_DATETIME + FIVE_SECONDS
rl_obj.current_count = 3
rl_obj._update_limit_period()
self.assertEqual(START_DATETIME + FIVE_SECONDS,
rl_obj.current_limit_period_start)
self.assertEqual(0, rl_obj.current_count)
def test_withdraw_rate_limiting(self, now_mock):
now_mock.return_value = START_DATETIME
item_limit = 10
rl_obj = rate_limiter.RateLimiter(item_limit, 5)
rl_obj.add_items(range(0, 105))
# We should've gotten the first ten items back.
return_list = rl_obj.withdraw_items()
self.assertEqual(list(range(0, 10)), return_list)
self.assertEqual(95, len(rl_obj))
# Now we should only get the empty list back.
for n in range(0, 3):
return_list = rl_obj.withdraw_items()
self.assertEqual([], return_list)
self.assertEqual(95, len(rl_obj))
item_count = 95
item_start = 10
item_end = 20
# Pump items out of the rate limiter.
for i in range(0, 9):
rl_obj._start_new_limit_period()
return_list = rl_obj.withdraw_items()
item_count -= item_limit
self.assertEqual(list(range(item_start, item_end)), return_list)
self.assertEqual(item_count, len(rl_obj))
item_start += item_limit
item_end += item_limit
# Now we should only get 5 back.
rl_obj._start_new_limit_period()
return_list = rl_obj.withdraw_items()
self.assertEqual(list(range(100, 105)), return_list)
self.assertEqual(0, len(rl_obj))
self.assertEqual(5, rl_obj.current_count)
def test_adding_items(self):
rl_obj = rate_limiter.RateLimiter(3, 10)
rl_obj.add_items(range(0, 10))
self.assertEqual(10, len(rl_obj))
rl_obj.add_items(range(0, 10))
self.assertEqual(20, len(rl_obj))