def test_global_max_tokens(self, mocked_now): mocked_now.return_value = fake_now ramp = rate_limiter.RateLimiter(global_max_tokens=499, ) self.assertEqual(ramp._phase, 0) self.assertEqual(ramp._maximum_tokens, 499) ramp.take_tokens() # Advance the clock 1 phase mocked_now.return_value = now_plus_n( seconds=rate_limiter.default_phase_length, microseconds=1, ) ramp.take_tokens() self.assertEqual(ramp._phase, 1) self.assertEqual(ramp._maximum_tokens, 499) # Advance the clock another phase mocked_now.return_value = now_plus_n( seconds=rate_limiter.default_phase_length * 2, microseconds=1, ) ramp.take_tokens() self.assertEqual(ramp._phase, 2) self.assertEqual(ramp._maximum_tokens, 499) # Advance the clock another ms and the phase should not advance mocked_now.return_value = now_plus_n( seconds=rate_limiter.default_phase_length * 2, microseconds=2, ) ramp.take_tokens() self.assertEqual(ramp._phase, 2) self.assertEqual(ramp._maximum_tokens, 499)
def test_take_batch_size(self, mocked_now): """Verifies that if the clock advances but nothing happens, the RateLimiter doesn't ramp up. """ page_size: int = 20 mocked_now.return_value = fake_now ramp = rate_limiter.RateLimiter() ramp._available_tokens = 15 self.assertEqual(ramp.take_tokens(page_size, allow_less=True), 15) # Advance the clock 1 phase mocked_now.return_value = now_plus_n( seconds=rate_limiter.default_phase_length, microseconds=1, ) ramp._check_phase() self.assertEqual(ramp._maximum_tokens, 750) for _ in range(740 // page_size): self.assertEqual( ramp.take_tokens(page_size), page_size, msg=f"page {_} should have been allowed", ) self.assertEqual(ramp.take_tokens(page_size, allow_less=True), 10) self.assertEqual(ramp.take_tokens(page_size, allow_less=True), 0)
def test_rate_limiter_basic(self, mocked_now): """Verifies that if the clock does not advance, the RateLimiter allows 500 writes before crashing out. """ mocked_now.return_value = fake_now # This RateLimiter will never advance. Poor fella. ramp = rate_limiter.RateLimiter() for _ in range(rate_limiter.default_initial_tokens): self.assertEqual(ramp.take_tokens(), 1) self.assertEqual(ramp.take_tokens(), 0)
def test_rate_limiter_basic(mocked_now): """Verifies that if the clock does not advance, the RateLimiter allows 500 writes before crashing out. """ from google.cloud.firestore_v1 import rate_limiter mocked_now.return_value = fake_now # This RateLimiter will never advance. Poor fella. ramp = rate_limiter.RateLimiter() for _ in range(rate_limiter.default_initial_tokens): assert ramp.take_tokens() == 1 assert ramp.take_tokens() == 0
def test_rate_limiter_with_refill(self, mocked_now): """Verifies that if the clock advances, the RateLimiter allows appropriate additional writes. """ mocked_now.return_value = fake_now ramp = rate_limiter.RateLimiter() ramp._available_tokens = 0 self.assertEqual(ramp.take_tokens(), 0) # Advance the clock 0.1 seconds mocked_now.return_value = now_plus_n(microseconds=100000) for _ in range(round(rate_limiter.default_initial_tokens / 10)): self.assertEqual(ramp.take_tokens(), 1) self.assertEqual(ramp.take_tokens(), 0)
def test_rate_limiter_with_refill(mocked_now): """Verifies that if the clock advances, the RateLimiter allows appropriate additional writes. """ from google.cloud.firestore_v1 import rate_limiter mocked_now.return_value = fake_now ramp = rate_limiter.RateLimiter() ramp._available_tokens = 0 assert ramp.take_tokens() == 0 # Advance the clock 0.1 seconds mocked_now.return_value = now_plus_n(microseconds=100000) for _ in range(round(rate_limiter.default_initial_tokens / 10)): assert ramp.take_tokens() == 1 assert ramp.take_tokens() == 0
def test_rate_limiter_phase_length(self, mocked_now): """Verifies that if the clock advances, the RateLimiter allows appropriate additional writes. """ mocked_now.return_value = fake_now ramp = rate_limiter.RateLimiter() self.assertEqual(ramp.take_tokens(), 1) ramp._available_tokens = 0 self.assertEqual(ramp.take_tokens(), 0) # Advance the clock 1 phase mocked_now.return_value = now_plus_n( seconds=rate_limiter.default_phase_length, microseconds=1, ) for _ in range(round(rate_limiter.default_initial_tokens * 3 / 2)): self.assertTrue(ramp.take_tokens(), msg=f"token {_} should have been allowed") self.assertEqual(ramp.take_tokens(), 0)
def test_rate_limiter_idle_phase_length(self, mocked_now): """Verifies that if the clock advances but nothing happens, the RateLimiter doesn't ramp up. """ mocked_now.return_value = fake_now ramp = rate_limiter.RateLimiter() ramp._available_tokens = 0 self.assertEqual(ramp.take_tokens(), 0) # Advance the clock 1 phase mocked_now.return_value = now_plus_n( seconds=rate_limiter.default_phase_length, microseconds=1, ) for _ in range(round(rate_limiter.default_initial_tokens)): self.assertEqual(ramp.take_tokens(), 1, msg=f"token {_} should have been allowed") self.assertEqual(ramp._maximum_tokens, 500) self.assertEqual(ramp.take_tokens(), 0)
def test_rate_limiter_idle_phase_length(mocked_now): """Verifies that if the clock advances but nothing happens, the RateLimiter doesn't ramp up. """ from google.cloud.firestore_v1 import rate_limiter mocked_now.return_value = fake_now ramp = rate_limiter.RateLimiter() ramp._available_tokens = 0 assert ramp.take_tokens() == 0 # Advance the clock 1 phase mocked_now.return_value = now_plus_n( seconds=rate_limiter.default_phase_length, microseconds=1, ) for _ in range(round(rate_limiter.default_initial_tokens)): assert ramp.take_tokens() == 1 assert ramp._maximum_tokens == 500 assert ramp.take_tokens() == 0
def test_rate_limiter_phase_length(mocked_now): """Verifies that if the clock advances, the RateLimiter allows appropriate additional writes. """ from google.cloud.firestore_v1 import rate_limiter mocked_now.return_value = fake_now ramp = rate_limiter.RateLimiter() assert ramp.take_tokens() == 1 ramp._available_tokens = 0 assert ramp.take_tokens() == 0 # Advance the clock 1 phase mocked_now.return_value = now_plus_n( seconds=rate_limiter.default_phase_length, microseconds=1, ) for _ in range(round(rate_limiter.default_initial_tokens * 3 / 2)): assert ramp.take_tokens() assert ramp.take_tokens() == 0
def test_take_batch_size(mocked_now): """Verifies that if the clock advances but nothing happens, the RateLimiter doesn't ramp up. """ from google.cloud.firestore_v1 import rate_limiter page_size: int = 20 mocked_now.return_value = fake_now ramp = rate_limiter.RateLimiter() ramp._available_tokens = 15 assert ramp.take_tokens(page_size, allow_less=True) == 15 # Advance the clock 1 phase mocked_now.return_value = now_plus_n( seconds=rate_limiter.default_phase_length, microseconds=1, ) ramp._check_phase() assert ramp._maximum_tokens == 750 for _ in range(740 // page_size): assert ramp.take_tokens(page_size) == page_size assert ramp.take_tokens(page_size, allow_less=True) == 10 assert ramp.take_tokens(page_size, allow_less=True) == 0
def test_global_max_tokens(mocked_now): from google.cloud.firestore_v1 import rate_limiter mocked_now.return_value = fake_now ramp = rate_limiter.RateLimiter(global_max_tokens=499, ) assert ramp._phase == 0 assert ramp._maximum_tokens == 499 ramp.take_tokens() # Advance the clock 1 phase mocked_now.return_value = now_plus_n( seconds=rate_limiter.default_phase_length, microseconds=1, ) ramp.take_tokens() assert ramp._phase == 1 assert ramp._maximum_tokens == 499 # Advance the clock another phase mocked_now.return_value = now_plus_n( seconds=rate_limiter.default_phase_length * 2, microseconds=1, ) ramp.take_tokens() assert ramp._phase == 2 assert ramp._maximum_tokens == 499 # Advance the clock another ms and the phase should not advance mocked_now.return_value = now_plus_n( seconds=rate_limiter.default_phase_length * 2, microseconds=2, ) ramp.take_tokens() assert ramp._phase == 2 assert ramp._maximum_tokens == 499