"""Tests for the Redis-backed token-bucket rate limiter. These tests use fakeredis (Lua enabled via lupa) so no live Redis server is required. """ import asyncio import sys import pytest from fastapi import Depends, FastAPI, Request from fastapi.testclient import TestClient sys.path.insert(0, "/workspace") import fakeredis.aioredis from redis.asyncio import Redis from redis.exceptions import ConnectionError from app.rate_limit import ( RateLimit, RateLimitConfig, RateLimiter, default_user_scope, named_limit, rate_limit, ) @pytest.fixture() def fake_redis(): return fakeredis.aioredis.FakeRedis(decode_responses=True) def make_request(client_ip: str = "127.0.0.1") -> Request: scope = { "type": "http", "method": "GET", "path": "/", "headers": [], "client": (client_ip, 12345), } return Request(scope) @pytest.mark.asyncio async def test_token_bucket_allows_burst_then_exhausts(fake_redis): limiter = RateLimiter( RateLimitConfig( limit=RateLimit(name="burst", capacity=3, refill_rate=0.1), ), redis=fake_redis, ) request = make_request() results = [] for _ in range(4): result = await limiter.check(request, user_id=42) results.append(result) assert [r.allowed for r in results] == [True, True, True, False] assert results[0].remaining == 2 assert results[1].remaining == 1 assert results[2].remaining == 0 assert results[3].retry_after > 0 @pytest.mark.asyncio async def test_retry_after_reflects_refill_rate(fake_redis): limiter = RateLimiter( RateLimitConfig( limit=RateLimit(name="slow", capacity=1, refill_rate=0.5), ), redis=fake_redis, ) request = make_request() assert (await limiter.check(request, user_id=7)).allowed is True denied = await limiter.check(request, user_id=7) assert denied.allowed is False # 1 token / 0.5 tokens per second = 2 seconds. assert denied.retry_after == 2 @pytest.mark.asyncio async def test_separate_scopes_get_separate_buckets(fake_redis): limiter = RateLimiter( RateLimitConfig( limit=RateLimit(name="scoped", capacity=1, refill_rate=0.1), ), redis=fake_redis, ) request = make_request() assert (await limiter.check(request, user_id=1)).allowed is True assert (await limiter.check(request, user_id=2)).allowed is True assert (await limiter.check(request, user_id=1)).allowed is False @pytest.mark.asyncio async def test_fail_open_returns_allowed_when_redis_unavailable(): class BrokenRedis(Redis): async def eval(self, *args, **kwargs): raise ConnectionError("down") limiter = RateLimiter( RateLimitConfig( limit=RateLimit(name="open", capacity=1, refill_rate=0.1, fail_open=True), ), redis=BrokenRedis(), ) result = await limiter.check(make_request(), user_id=1) assert result.allowed is True @pytest.mark.asyncio async def test_fail_closed_returns_denied_with_negative_retry_after(): class BrokenRedis(Redis): async def eval(self, *args, **kwargs): raise ConnectionError("down") limiter = RateLimiter( RateLimitConfig( limit=RateLimit(name="closed", capacity=1, refill_rate=0.1, fail_open=False), ), redis=BrokenRedis(), ) result = await limiter.check(make_request(), user_id=1) assert result.allowed is False assert result.retry_after == -1 @pytest.mark.asyncio async def test_fail_closed_call_raises_503(fake_redis): class BrokenRedis(Redis): async def eval(self, *args, **kwargs): raise ConnectionError("down") limiter = RateLimiter( RateLimitConfig( limit=RateLimit(name="closed", capacity=1, refill_rate=0.1, fail_open=False), ), redis=BrokenRedis(), ) with pytest.raises(Exception) as exc_info: await limiter(make_request(), user_id=1) assert exc_info.value.status_code == 503 def test_named_limit_builds_expected_config(): config = named_limit( "chat_send", capacity=5, refill_rate=0.2, fail_open=False, ) assert config.limit.name == "chat_send" assert config.limit.capacity == 5 assert config.limit.refill_rate == 0.2 assert config.limit.fail_open is False assert config.scope is default_user_scope def test_invalid_limit_definition_rejected(): with pytest.raises(ValueError): named_limit("bad", capacity=0, refill_rate=1.0) with pytest.raises(ValueError): named_limit("bad", capacity=1, refill_rate=0) def test_dependency_integration_returns_429_with_retry_after(): async def fake_user_id() -> int: return 99 chat_limit = rate_limit( named_limit("chat_send_test", capacity=2, refill_rate=0.1, fail_open=False), user_id_dependency=Depends(fake_user_id), ) app = FastAPI() @app.get("/probe") async def probe(_: None = Depends(chat_limit)): return {"ok": True} client = TestClient(app) assert client.get("/probe").status_code == 200 assert client.get("/probe").status_code == 200 denied = client.get("/probe") assert denied.status_code == 429 assert denied.headers["retry-after"] assert denied.json() == {"detail": "Rate limit exceeded"}