def test_moving_window_redis(self): storage = RedisStorage("redis://localhost:6379") limiter = MovingWindowRateLimiter(storage) limit = RateLimitItemPerSecond(10, 2) for i in range(0,10): self.assertTrue(limiter.hit(limit)) self.assertEqual(limiter.get_window_stats(limit)[1], 10 - (i + 1)) time.sleep(2*0.095) self.assertFalse(limiter.hit(limit)) time.sleep(0.4) self.assertTrue(limiter.hit(limit)) self.assertTrue(limiter.hit(limit)) self.assertEqual(limiter.get_window_stats(limit)[1], 0)
def test_moving_window_redis(self): storage = RedisStorage("redis://localhost:6379") limiter = MovingWindowRateLimiter(storage) limit = RateLimitItemPerSecond(10, 2) for i in range(0, 10): self.assertTrue(limiter.hit(limit)) self.assertEqual(limiter.get_window_stats(limit)[1], 10 - (i + 1)) time.sleep(2 * 0.095) self.assertFalse(limiter.hit(limit)) time.sleep(0.4) self.assertTrue(limiter.hit(limit)) self.assertTrue(limiter.hit(limit)) self.assertEqual(limiter.get_window_stats(limit)[1], 0)
class RateLimiter: def __init__(self, storage, limit, *, identifiers=None, metrics): if identifiers is None: identifiers = [] self._window = MovingWindowRateLimiter(storage) self._limits = parse_many(limit) self._identifiers = identifiers self._metrics = metrics def _get_identifiers(self, identifiers): return [str(i) for i in list(self._identifiers) + list(identifiers)] @_return_on_exception(True, redis.RedisError) def test(self, *identifiers): return all( [ self._window.test(limit, *self._get_identifiers(identifiers)) for limit in self._limits ] ) @_return_on_exception(True, redis.RedisError) def hit(self, *identifiers): return all( [ self._window.hit(limit, *self._get_identifiers(identifiers)) for limit in self._limits ] ) @_return_on_exception(None, redis.RedisError) def resets_in(self, *identifiers): resets = [] for limit in self._limits: resets_at, remaining = self._window.get_window_stats( limit, *self._get_identifiers(identifiers) ) # If this limit has any remaining limits left, then we will skip it # since it doesn't need reset. if remaining > 0: continue current = datetime.now(tz=timezone.utc) reset = datetime.fromtimestamp(resets_at, tz=timezone.utc) # If our current datetime is either greater than or equal to when # the limit resets, then we will skipp it since it has either # already reset, or it is resetting now. if current >= reset: continue # Add a timedelta that represents how long until this limit resets. resets.append(reset - current) # If we have any resets, then we'll go through and find whichever one # is going to reset soonest and use that as our hint for when this # limit might be available again. return first(sorted(resets))
class RateLimiter: def __init__(self, storage, limit, identifiers=None): if identifiers is None: identifiers = [] self._window = MovingWindowRateLimiter(storage) self._limits = parse_many(limit) self._identifiers = identifiers def _get_identifiers(self, identifiers): return [str(i) for i in list(self._identifiers) + list(identifiers)] def test(self, *identifiers): return all( [ self._window.test(limit, *self._get_identifiers(identifiers)) for limit in self._limits ] ) def hit(self, *identifiers): return all( [ self._window.hit(limit, *self._get_identifiers(identifiers)) for limit in self._limits ] ) def resets_in(self, *identifiers): resets = [] for limit in self._limits: resets_at, remaining = self._window.get_window_stats( limit, *self._get_identifiers(identifiers) ) # If this limit has any remaining limits left, then we will skip it # since it doesn't need reset. if remaining > 0: continue current = datetime.now(tz=timezone.utc) reset = datetime.fromtimestamp(resets_at, tz=timezone.utc) # If our current datetime is either greater than or equal to when # the limit resets, then we will skipp it since it has either # already reset, or it is resetting now. if current >= reset: continue # Add a timedelta that represents how long until this limit resets. resets.append(reset - current) # If we have any resets, then we'll go through and find whichever one # is going to reset soonest and use that as our hint for when this # limit might be available again. return first(sorted(resets))
def test_moving_window_in_memory(self): storage = MemoryStorage() limiter = MovingWindowRateLimiter(storage) with hiro.Timeline().freeze() as timeline: limit = RateLimitItemPerMinute(10) for i in range(0, 5): self.assertTrue(limiter.hit(limit)) self.assertTrue(limiter.hit(limit)) self.assertEqual( limiter.get_window_stats(limit)[1], 10 - ((i + 1) * 2)) timeline.forward(10) self.assertEqual(limiter.get_window_stats(limit)[1], 0) self.assertFalse(limiter.hit(limit)) timeline.forward(20) self.assertEqual(limiter.get_window_stats(limit)[1], 2) self.assertEqual( limiter.get_window_stats(limit)[0], int(time.time() + 30)) timeline.forward(31) self.assertEqual(limiter.get_window_stats(limit)[1], 10)
def test_moving_window_in_memory(self): storage = MemoryStorage() limiter = MovingWindowRateLimiter(storage) with hiro.Timeline().freeze() as timeline: limit = RateLimitItemPerMinute(10) for i in range(0,5): self.assertTrue(limiter.hit(limit)) self.assertTrue(limiter.hit(limit)) self.assertEqual( limiter.get_window_stats(limit)[1], 10 - ((i + 1) * 2) ) timeline.forward(10) self.assertEqual(limiter.get_window_stats(limit)[1], 0) self.assertFalse(limiter.hit(limit)) timeline.forward(20) self.assertEqual(limiter.get_window_stats(limit)[1], 2) self.assertEqual(limiter.get_window_stats(limit)[0], int(time.time() + 30)) timeline.forward(31) self.assertEqual(limiter.get_window_stats(limit)[1], 10)