"""Token-bucket rate limiter behaviour.""" from __future__ import annotations import time import pytest from media_server.services import rate_limit @pytest.fixture(autouse=True) def _reset_state(): rate_limit._state.clear() rate_limit._LAST_CLEANUP = 0.0 yield rate_limit._state.clear() def test_allows_up_to_capacity_then_blocks(monkeypatch): """Default execute bucket = 10/min.""" peer = "10.0.0.1" for i in range(10): ok, retry = rate_limit.check("execute", peer) assert ok, f"expected allow on attempt {i + 1}, got block (retry={retry})" ok, retry = rate_limit.check("execute", peer) assert ok is False assert retry is not None and retry > 0 def test_different_peers_independent(): for _ in range(10): assert rate_limit.check("execute", "10.0.0.1")[0] # Different peer should still be allowed. assert rate_limit.check("execute", "10.0.0.2")[0] def test_unknown_bucket_uses_default(): peer = "10.0.0.3" # default = 60/min — first call always allowed. allowed, _ = rate_limit.check("nonexistent-bucket", peer) assert allowed def test_auth_bucket_is_strict(): """auth bucket = 5/min.""" peer = "10.0.0.4" for _ in range(5): assert rate_limit.check("auth", peer)[0] blocked, retry = rate_limit.check("auth", peer) assert not blocked assert retry is not None def test_refill_eventually_unblocks(monkeypatch): """Verify the bucket refills — exhaust then wait one refill period.""" peer = "10.0.0.5" # Replace BUCKETS with a fast-refilling one for the test only. monkeypatch.setitem( rate_limit.BUCKETS, "fast-test", rate_limit.BucketConfig(capacity=2, refill_per_sec=10.0), ) assert rate_limit.check("fast-test", peer)[0] assert rate_limit.check("fast-test", peer)[0] assert not rate_limit.check("fast-test", peer)[0] time.sleep(0.15) # 0.15 * 10 = 1.5 tokens assert rate_limit.check("fast-test", peer)[0]