Esempio n. 1
0
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)
Esempio n. 2
0
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
Esempio n. 4
0
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)
Esempio n. 5
0
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)