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))