daddy
de2d431571
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
125 lines
3.8 KiB
Python
125 lines
3.8 KiB
Python
"""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
|
|
from app.config import settings
|
|
from app.main import app
|
|
|
|
|
|
@pytest.fixture(autouse=True)
|
|
def _clean_limiter():
|
|
# This module tests the limiter itself: force it on (the checkpoint
|
|
# runner disables it suite-wide via RATE_LIMIT_ENABLED) and start each
|
|
# test from an empty window.
|
|
prev = settings.RATE_LIMIT_ENABLED
|
|
settings.RATE_LIMIT_ENABLED = True
|
|
rate_limit.reset_all()
|
|
yield
|
|
rate_limit.reset_all()
|
|
settings.RATE_LIMIT_ENABLED = prev
|
|
|
|
|
|
def test_login_is_limited_after_the_threshold():
|
|
client = TestClient(app)
|
|
codes = []
|
|
for i in range(15):
|
|
r = client.post("/api/auth/login", data={"username": f"nobody{i}", "password": "x"})
|
|
codes.append(r.status_code)
|
|
# First 10 get through to the (failing) auth check; the rest are 429.
|
|
assert codes[:10] == [401] * 10
|
|
assert 429 in codes[10:]
|
|
assert codes[-1] == 429
|
|
|
|
|
|
def test_register_shares_the_same_bucket():
|
|
client = TestClient(app)
|
|
# Burn the budget on login...
|
|
for i in range(10):
|
|
client.post("/api/auth/login", data={"username": f"x{i}", "password": "y"})
|
|
# ...register from the same IP is now blocked too.
|
|
r = client.post("/api/auth/register", json={
|
|
"username": "rl_user", "email": "rl@e.com",
|
|
"password": "secret123", "date_of_birth": "1990-01-01",
|
|
})
|
|
assert r.status_code == 429
|
|
assert "Retry-After" in r.headers
|
|
|
|
|
|
def test_limiter_window_is_per_key():
|
|
limiter = rate_limit.SlidingWindowLimiter(max_hits=2, window_seconds=60)
|
|
limiter.check("a")
|
|
limiter.check("a")
|
|
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
|