from __future__ import annotations import time from math import ceil import pytest from limits.aio.strategies import ( FixedWindowElasticExpiryRateLimiter, FixedWindowRateLimiter, MovingWindowRateLimiter, SlidingWindowCounterRateLimiter, ) from limits.limits import ( RateLimitItemPerHour, RateLimitItemPerMinute, RateLimitItemPerSecond, ) from limits.storage import storage_from_string from limits.storage.base import TimestampedSlidingWindow from tests.utils import ( async_all_storage, async_fixed_start, async_moving_window_storage, async_sliding_window_counter_storage, async_window, timestamp_based_key_ttl, ) @pytest.mark.asyncio @async_all_storage class TestAsyncFixedWindow: @async_fixed_start async def test_fixed_window(self, uri, args, fixture): storage = storage_from_string(uri, **args) limiter = FixedWindowRateLimiter(storage) limit = RateLimitItemPerSecond(10, 2) async with async_window(1) as (start, _): assert all([await limiter.hit(limit) for _ in range(0, 10)]) assert not await limiter.hit(limit) assert (await limiter.get_window_stats(limit)).remaining == 0 assert (await limiter.get_window_stats(limit)).reset_time == pytest.approx( start + 2, 1e-2 ) @async_fixed_start async def test_fixed_window_empty_stats(self, uri, args, fixture): storage = storage_from_string(uri, **args) limiter = FixedWindowRateLimiter(storage) limit = RateLimitItemPerSecond(10, 2) assert (await limiter.get_window_stats(limit)).remaining == 10 assert (await limiter.get_window_stats(limit)).reset_time == pytest.approx( time.time(), 1e-2 ) @async_fixed_start async def test_fixed_window_multiple_cost(self, uri, args, fixture): storage = storage_from_string(uri, **args) limiter = FixedWindowRateLimiter(storage) limit = RateLimitItemPerMinute(10, 2) assert not await limiter.hit(limit, "k1", cost=11) assert await limiter.hit(limit, "k2", cost=5) assert (await limiter.get_window_stats(limit, "k2")).remaining == 5 assert not await limiter.test(limit, "k2", cost=6) assert not await limiter.hit(limit, "k2", cost=6) @async_fixed_start async def test_fixed_window_with_elastic_expiry(self, uri, args, fixture): storage = storage_from_string(uri, **args) with pytest.warns(DeprecationWarning): limiter = FixedWindowElasticExpiryRateLimiter(storage) limit = RateLimitItemPerSecond(10, 2) async with async_window(1) as (start, end): assert all([await limiter.hit(limit) for _ in range(0, 10)]) assert not await limiter.hit(limit) assert (await limiter.get_window_stats(limit)).remaining == 0 assert (await limiter.get_window_stats(limit)).reset_time == pytest.approx( start + 2, 1e-2 ) async with async_window(3) as (start, end): assert not await limiter.hit(limit) assert await limiter.hit(limit) assert (await limiter.get_window_stats(limit)).remaining == 9 assert (await limiter.get_window_stats(limit)).reset_time == pytest.approx( end + 2, 1e-2 ) @async_fixed_start async def test_fixed_window_with_elastic_expiry_multiple_cost( self, uri, args, fixture ): storage = storage_from_string(uri, **args) with pytest.warns(DeprecationWarning): limiter = FixedWindowElasticExpiryRateLimiter(storage) limit = RateLimitItemPerSecond(10, 2) assert not await limiter.hit(limit, "k1", cost=11) async with async_window(0) as (_, end): assert await limiter.hit(limit, "k2", cost=5) assert (await limiter.get_window_stats(limit, "k2")).remaining == 5 assert ( await limiter.get_window_stats(limit, "k2") ).reset_time == pytest.approx(end + 2, 1e-2) assert not await limiter.hit(limit, "k2", cost=6) @async_fixed_start @pytest.mark.flaky async def test_test_fixed_window(self, uri, args, fixture): storage = storage_from_string(uri, **args) limiter = FixedWindowRateLimiter(storage) limit = RateLimitItemPerHour(2, 1) assert await limiter.hit(limit) assert await limiter.test(limit) assert await limiter.hit(limit) assert not await limiter.test(limit) assert not await limiter.hit(limit) @pytest.mark.asyncio @async_moving_window_storage class TestAsyncMovingWindow: async def test_moving_window_stats(self, uri, args, fixture): storage = storage_from_string(uri, **args) limiter = MovingWindowRateLimiter(storage) limit = RateLimitItemPerMinute(2) assert await limiter.hit(limit, "key") time.sleep(1) assert await limiter.hit(limit, "key") time.sleep(1) assert not await limiter.hit(limit, "key") assert (await limiter.get_window_stats(limit, "key")).remaining == 0 assert ( await limiter.get_window_stats(limit, "key") ).reset_time - time.time() == pytest.approx(58, 1e-2) async def test_moving_window(self, uri, args, fixture): storage = storage_from_string(uri, **args) limiter = MovingWindowRateLimiter(storage) limit = RateLimitItemPerSecond(10, 2) # 5 hits in the first 500ms async with async_window(0.5): assert all([await limiter.hit(limit) for i in range(5)]) # 5 hits in the last 200ms async with async_window(2, delay=1.3): assert all([await limiter.hit(limit) for i in range(5)]) # 11th fails assert not await limiter.hit(limit) # 5 more succeed since there were only 5 in the last 2 seconds assert all([await limiter.hit(limit) for i in range(5)]) assert (await limiter.get_window_stats(limit)).remaining == 0 async def test_moving_window_empty_stats(self, uri, args, fixture): storage = storage_from_string(uri, **args) limiter = MovingWindowRateLimiter(storage) limit = RateLimitItemPerSecond(10, 2) assert (await limiter.get_window_stats(limit)).remaining == 10 assert (await limiter.get_window_stats(limit)).reset_time == pytest.approx( time.time() + 2, 1e-2 ) async def test_moving_window_multiple_cost(self, uri, args, fixture): storage = storage_from_string(uri, **args) limiter = MovingWindowRateLimiter(storage) limit = RateLimitItemPerSecond(10, 2) assert not await limiter.hit(limit, "k1", cost=11) # 5 hits in the first 100ms async with async_window(0.1): assert await limiter.hit(limit, "k2", cost=5) # 5 hits in the last 100ms async with async_window(2, delay=1.8): assert all([await limiter.hit(limit, "k2") for i in range(4)]) assert not await limiter.test(limit, "k2", cost=2) assert not await limiter.hit(limit, "k2", cost=2) assert await limiter.hit(limit, "k2") assert all([await limiter.hit(limit, "k2") for i in range(5)]) assert (await limiter.get_window_stats(limit, "k2")).remaining == 0 assert not await limiter.hit(limit, "k2", cost=2) async def test_moving_window_varying_cost(self, uri, args, fixture): storage = storage_from_string(uri, **args) limiter = MovingWindowRateLimiter(storage) five_per_min = RateLimitItemPerMinute(5) await limiter.hit(five_per_min, cost=5) assert not await limiter.hit(five_per_min, cost=2) await limiter.clear(five_per_min) assert await limiter.hit(five_per_min) async def test_moving_window_huge_cost_async(self, uri, args, fixture): storage = storage_from_string(uri, **args) limiter = MovingWindowRateLimiter(storage) many_per_min = RateLimitItemPerMinute(1_000_000) await limiter.hit(many_per_min, cost=999_999) assert not await limiter.hit(many_per_min, cost=2) await limiter.clear(many_per_min) assert await limiter.hit(many_per_min) async def test_test_moving_window(self, uri, args, fixture): storage = storage_from_string(uri, **args) limit = RateLimitItemPerHour(2, 1) limiter = MovingWindowRateLimiter(storage) assert await limiter.hit(limit) assert await limiter.test(limit) assert await limiter.hit(limit) assert not await limiter.test(limit) assert not await limiter.hit(limit) @pytest.mark.asyncio @async_sliding_window_counter_storage class TestAsyncSlidingWindow: @async_fixed_start async def test_sliding_window_counter(self, uri, args, fixture): storage = storage_from_string(uri, **args) limiter = SlidingWindowCounterRateLimiter(storage) limit = RateLimitItemPerSecond(10, 2) if isinstance(storage, TimestampedSlidingWindow): # Avoid testing the behaviour when the window is about to be reset ttl = timestamp_based_key_ttl(limit) if ttl < 1: time.sleep(ttl) async with async_window(1) as (start, _): assert all([await limiter.hit(limit) for _ in range(0, 10)]) assert not await limiter.hit(limit) assert (await limiter.get_window_stats(limit)).remaining == 0 assert (await limiter.get_window_stats(limit)).reset_time == pytest.approx( start + 2, 1e-2 ) @pytest.mark.flaky async def test_sliding_window_counter_total_reset(self, uri, args, fixture): storage = storage_from_string(uri, **args) limiter = SlidingWindowCounterRateLimiter(storage) multiple = 10 period = 1 limit = RateLimitItemPerSecond(multiple, period) assert (await limiter.get_window_stats(limit)).remaining == multiple if isinstance(storage, TimestampedSlidingWindow): # Avoid testing the behaviour when the window is about to be reset ttl = timestamp_based_key_ttl(limit) if ttl < 0.5: time.sleep(ttl) assert await limiter.hit(limit, cost=multiple) assert not await limiter.hit(limit) assert (await limiter.get_window_stats(limit)).remaining == 0 time.sleep(period * 2) assert (await limiter.get_window_stats(limit)).remaining == multiple assert (await limiter.get_window_stats(limit)).reset_time == pytest.approx( time.time(), abs=1e-2 ) async def test_sliding_window_counter_current_window(self, uri, args, fixture): """Check the window stats when only the current window is filled""" storage = storage_from_string(uri, **args) limiter = SlidingWindowCounterRateLimiter(storage) limit = RateLimitItemPerHour(2, 24) if isinstance(storage, TimestampedSlidingWindow): # Avoid testing the behaviour when the window is about to be reset ttl = timestamp_based_key_ttl(limit) if ttl < 0.5: time.sleep(ttl) assert await limiter.hit(limit) now = time.time() if isinstance(storage, TimestampedSlidingWindow): expected_reset_time = now + timestamp_based_key_ttl(limit, now) else: expected_reset_time = now + 24 * 3600 assert (await limiter.get_window_stats(limit)).reset_time == pytest.approx( expected_reset_time, 1e-2 ) assert (await limiter.get_window_stats(limit)).remaining == 1 assert await limiter.hit(limit) assert not await limiter.hit(limit) @pytest.mark.flaky(max_runs=3) async def test_sliding_window_counter_previous_window(self, uri, args, fixture): """Check the window stats when the previous window is partially filled""" storage = storage_from_string(uri, **args) limiter = SlidingWindowCounterRateLimiter(storage) limit = RateLimitItemPerSecond(5, 1) sleep_margin = 0.001 if isinstance(storage, TimestampedSlidingWindow): # Avoid testing the behaviour when the window is about to be reset ttl = timestamp_based_key_ttl(limit) if ttl < 0.3: time.sleep(ttl + sleep_margin) t0 = time.time() previous_window_hits = 3 await limiter.hit(limit) t1 = time.time() for i in range(previous_window_hits - 1): await limiter.hit(limit) # Check the stats: only the current window is filled if isinstance(storage, TimestampedSlidingWindow): expected_reset_time = t0 + timestamp_based_key_ttl(limit, t0) else: expected_reset_time = t1 + 1 reset_time = (await limiter.get_window_stats(limit)).reset_time assert reset_time == pytest.approx(expected_reset_time, abs=0.03) assert (await limiter.get_window_stats(limit)).remaining == 2 # Wait for the next window sleep_time = expected_reset_time - time.time() + sleep_margin time.sleep(sleep_time) # A new hit should be available immediately after window shift # The limiter should reset in a fraction of a period, according to how many hits are in the previous window reset_time = (await limiter.get_window_stats(limit)).reset_time reset_in = reset_time - time.time() assert reset_in == pytest.approx( limit.get_expiry() / previous_window_hits, abs=0.03 ) assert (await limiter.get_window_stats(limit)).remaining == 3 assert await limiter.hit(limit) assert await limiter.hit(limit) for i in range(previous_window_hits): # A new item hit should be freed by the previous window t0 = time.time() assert (await limiter.get_window_stats(limit)).remaining == 1 assert await limiter.hit(limit) assert (await limiter.get_window_stats(limit)).remaining == 0 assert not await limiter.hit(limit) # The previous window has 4 hits. The reset time should be in a 1/4 of the window expiry reset_time = (await limiter.get_window_stats(limit)).reset_time t1 = time.time() reset_in = reset_time - time.time() assert reset_in == pytest.approx( limit.get_expiry() / previous_window_hits - (t1 - t0), abs=0.03 ) # Wait for the next hit available time.sleep(reset_in + sleep_margin) @async_fixed_start async def test_sliding_window_counter_empty_stats(self, uri, args, fixture): storage = storage_from_string(uri, **args) limiter = SlidingWindowCounterRateLimiter(storage) limit = RateLimitItemPerSecond(10, 2) assert (await limiter.get_window_stats(limit)).remaining == 10 assert (await limiter.get_window_stats(limit)).reset_time == pytest.approx( time.time(), 1e-2 ) @async_fixed_start @pytest.mark.flaky async def test_sliding_window_counter_stats(self, uri, args, fixture): storage = storage_from_string(uri, **args) limiter = SlidingWindowCounterRateLimiter(storage) limit = RateLimitItemPerMinute(2) if isinstance(storage, TimestampedSlidingWindow): next_second_from_now = ceil(time.time()) assert await limiter.hit(limit, "key") time.sleep(1) assert await limiter.hit(limit, "key") time.sleep(1) assert not await limiter.hit(limit, "key") assert (await limiter.get_window_stats(limit, "key")).remaining == 0 if isinstance(storage, TimestampedSlidingWindow): # With timestamp-based key implementation, # the reset time is periodic according to the worker's timestamp reset_time = (await limiter.get_window_stats(limit, "key")).reset_time expected_reset = int( limit.get_expiry() - (next_second_from_now % limit.get_expiry()) ) assert reset_time - next_second_from_now == pytest.approx( expected_reset, abs=1e-2 ) else: assert ( await limiter.get_window_stats(limit, "key") ).reset_time - time.time() == pytest.approx(58, 1e-2) @async_fixed_start async def test_sliding_window_counter_multiple_cost(self, uri, args, fixture): storage = storage_from_string(uri, **args) limiter = SlidingWindowCounterRateLimiter(storage) limit = RateLimitItemPerMinute(10, 2) if isinstance(storage, TimestampedSlidingWindow): # Avoid testing the behaviour when the window is about to be reset ttl = timestamp_based_key_ttl(limit) if ttl < 0.5: time.sleep(ttl) assert not await limiter.hit(limit, "k1", cost=11) assert await limiter.hit(limit, "k2", cost=5) assert (await limiter.get_window_stats(limit, "k2")).remaining == 5 assert not await limiter.test(limit, "k2", cost=6) assert not await limiter.hit(limit, "k2", cost=6) async def test_test_sliding_window_counter(self, uri, args, fixture): storage = storage_from_string(uri, **args) limit = RateLimitItemPerHour(2, 1) limiter = SlidingWindowCounterRateLimiter(storage) assert await limiter.hit(limit) assert await limiter.test(limit) assert await limiter.hit(limit) assert not await limiter.test(limit) assert not await limiter.hit(limit)