Generalize rate_limit.py into a reusable named-limit dependency
Adds rate_limit(name, max_hits=, window_seconds=) as a factory for FastAPI dependencies, generalizing the auth_rate_limit pattern for future endpoints (chat send, message report) without touching auth_rate_limit's own behavior. Reuses SlidingWindowLimiter for counting and the existing 429 + Retry-After response, and respects RATE_LIMIT_ENABLED the same way auth_rate_limit does. reset_all() now clears every registered limiter (auth's plus any named ones), so the conftest autouse fixture keeps working for sibling tasks automatically. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01Y2FESCaMXeJpC9gSctiEeD
This commit is contained in:
@@ -1,6 +1,7 @@
|
||||
"""The auth endpoints are rate-limited per client IP."""
|
||||
|
||||
import pytest
|
||||
from fastapi import Depends, FastAPI
|
||||
from fastapi.testclient import TestClient
|
||||
|
||||
from app import rate_limit
|
||||
@@ -54,3 +55,70 @@ def test_limiter_window_is_per_key():
|
||||
with pytest.raises(Exception):
|
||||
limiter.check("a")
|
||||
limiter.check("b") # different key, unaffected
|
||||
|
||||
|
||||
def _make_named_limit_app():
|
||||
"""A tiny standalone app exercising two independent named limits, so
|
||||
these tests don't depend on any sibling endpoint existing yet."""
|
||||
limit_a = rate_limit.rate_limit("test_named_a", max_hits=2, window_seconds=60)
|
||||
limit_b = rate_limit.rate_limit("test_named_b", max_hits=2, window_seconds=60)
|
||||
|
||||
app = FastAPI()
|
||||
|
||||
@app.post("/a", dependencies=[Depends(limit_a)])
|
||||
def _a():
|
||||
return {"ok": True}
|
||||
|
||||
@app.post("/b", dependencies=[Depends(limit_b)])
|
||||
def _b():
|
||||
return {"ok": True}
|
||||
|
||||
return TestClient(app)
|
||||
|
||||
|
||||
def test_named_limits_do_not_share_counters():
|
||||
client = _make_named_limit_app()
|
||||
|
||||
# Burn out limit "a" entirely.
|
||||
assert client.post("/a").status_code == 200
|
||||
assert client.post("/a").status_code == 200
|
||||
r = client.post("/a")
|
||||
assert r.status_code == 429
|
||||
assert "Retry-After" in r.headers
|
||||
|
||||
# Limit "b" is untouched -- same process, same client IP, different name.
|
||||
assert client.post("/b").status_code == 200
|
||||
assert client.post("/b").status_code == 200
|
||||
assert client.post("/b").status_code == 429
|
||||
|
||||
|
||||
def test_named_limit_returns_429_with_retry_after_once_exhausted():
|
||||
client = _make_named_limit_app()
|
||||
client.post("/a")
|
||||
client.post("/a")
|
||||
r = client.post("/a")
|
||||
assert r.status_code == 429
|
||||
retry_after = r.headers["Retry-After"]
|
||||
assert int(retry_after) >= 1
|
||||
|
||||
|
||||
def test_named_limit_is_a_noop_when_rate_limiting_disabled():
|
||||
prev = settings.RATE_LIMIT_ENABLED
|
||||
settings.RATE_LIMIT_ENABLED = False
|
||||
try:
|
||||
client = _make_named_limit_app()
|
||||
codes = [client.post("/a").status_code for _ in range(10)]
|
||||
assert codes == [200] * 10
|
||||
finally:
|
||||
settings.RATE_LIMIT_ENABLED = prev
|
||||
|
||||
|
||||
def test_reset_all_clears_named_limiters_too():
|
||||
client = _make_named_limit_app()
|
||||
client.post("/a")
|
||||
client.post("/a")
|
||||
assert client.post("/a").status_code == 429
|
||||
|
||||
rate_limit.reset_all()
|
||||
|
||||
assert client.post("/a").status_code == 200
|
||||
|
||||
Reference in New Issue
Block a user