modbison/backend/tests/test_rate_limit.py

"""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"}