generated from kgod/ai-review-template
91 lines
2.8 KiB
Python
91 lines
2.8 KiB
Python
"""Per-session sliding-window rate limiting for light optimization."""
|
|
|
|
from fastapi import FastAPI
|
|
from fastapi.testclient import TestClient
|
|
from fastapi.responses import JSONResponse
|
|
|
|
from app.fsm import FSMError
|
|
from app.optimization_models import OptimizationRunView
|
|
from app.rate_limit import SlidingWindowRateLimiter
|
|
from app.resume_routes import register_resume_routes
|
|
|
|
|
|
class FakeClock:
|
|
def __init__(self) -> None:
|
|
self.now = 1000.0
|
|
|
|
def __call__(self) -> float:
|
|
return self.now
|
|
|
|
|
|
def test_allows_up_to_limit_then_rejects() -> None:
|
|
clock = FakeClock()
|
|
limiter = SlidingWindowRateLimiter(limit=3, window_seconds=3600, clock=clock)
|
|
|
|
assert limiter.allow("s1") is True
|
|
assert limiter.allow("s1") is True
|
|
assert limiter.allow("s1") is True
|
|
assert limiter.allow("s1") is False
|
|
|
|
|
|
def test_window_slides_and_allows_again() -> None:
|
|
clock = FakeClock()
|
|
limiter = SlidingWindowRateLimiter(limit=2, window_seconds=3600, clock=clock)
|
|
|
|
assert limiter.allow("s1") is True
|
|
assert limiter.allow("s1") is True
|
|
assert limiter.allow("s1") is False
|
|
|
|
clock.now += 3601
|
|
assert limiter.allow("s1") is True
|
|
|
|
|
|
def test_sessions_are_independent() -> None:
|
|
limiter = SlidingWindowRateLimiter(
|
|
limit=1, window_seconds=3600, clock=FakeClock()
|
|
)
|
|
|
|
assert limiter.allow("s1") is True
|
|
assert limiter.allow("s2") is True
|
|
assert limiter.allow("s1") is False
|
|
|
|
|
|
def test_light_route_returns_429_without_calling_optimizer_when_limited() -> None:
|
|
view = OptimizationRunView(
|
|
id="r1",
|
|
mode="light",
|
|
status="proposal_pending",
|
|
entry_id="e1",
|
|
)
|
|
|
|
class StubAgent:
|
|
def __init__(self) -> None:
|
|
self.calls = 0
|
|
|
|
def optimize_light(
|
|
self, session_id: str, request: object
|
|
) -> OptimizationRunView:
|
|
self.calls += 1
|
|
return view
|
|
|
|
application = FastAPI()
|
|
|
|
@application.exception_handler(FSMError)
|
|
async def handle_fsm_error(_request: object, exc: FSMError) -> JSONResponse:
|
|
return JSONResponse(status_code=exc.status_code, content={"detail": exc.message})
|
|
application.state.light_opt_limiter = SlidingWindowRateLimiter(
|
|
limit=2, window_seconds=3600
|
|
)
|
|
agent = StubAgent()
|
|
register_resume_routes(application, agent, "/ai-api/resume-agent")
|
|
client = TestClient(application, raise_server_exceptions=False)
|
|
url = "/ai-api/resume-agent/sessions/s1/resume/optimize/light"
|
|
|
|
assert client.post(url, json={"entry_id": "e1"}).status_code == 200
|
|
assert client.post(url, json={"entry_id": "e1"}).status_code == 200
|
|
blocked = client.post(url, json={"entry_id": "e1"})
|
|
|
|
assert blocked.status_code == 429
|
|
assert blocked.json()["detail"] == "操作过于频繁，请稍后再试（轻度优化每小时最多 20 次）。"
|
|
assert agent.calls == 2
|