from fastapi import FastAPI
from fastapi.testclient import TestClient


def test_rate_limit_rejects_after_limit():
    from axon.middleware.rate_limit import RateLimitMiddleware

    app = FastAPI()

    @app.get("/data")
    async def data():
        return {"ok": True}

    app.add_middleware(RateLimitMiddleware, limit_per_minute=1)
    client = TestClient(app)

    assert client.get("/data").status_code == 200
    assert client.get("/data").status_code == 429


def test_rate_limit_exempts_health():
    from axon.middleware.rate_limit import RateLimitMiddleware

    app = FastAPI()

    @app.get("/health")
    async def health():
        return {"ok": True}

    app.add_middleware(RateLimitMiddleware, limit_per_minute=1)
    client = TestClient(app)

    assert client.get("/health").status_code == 200
    assert client.get("/health").status_code == 200