177 lines
6.5 KiB
Python
177 lines
6.5 KiB
Python
"""限流闸门契约测试(文档 §3.5 的 429、§3.6 的 `RATE_LIMITED`,标注为可重试)。
|
|||
|
|
|
||
|
|
覆盖用户要求的三条:正常放行、超限拒绝且码与 `Retry-After` 正确、Redis 不可用时
|
||
|
|
降级放行;另加三条边界:
|
||
|
|
|
||
|
|
- 闸门关闭(`RATE_LIMIT_ENABLED=false`)时不访问后端;
|
||
|
|
- 认证先于限流(未带令牌仍是 401,不会被限流改写成 429);
|
||
|
|
- 计数维度按"用户 + 路由模板",而不是原始 URL——否则 `/{run_id}` 这类路径每个资源
|
||
|
|
都会各自计数,限流形同虚设。
|
||
|
|
|
||
|
|
全部进程内调用,不连 Redis(后端是替身),也不连数据库(放行路径用替身 Service)。
|
||
|
|
"""
|
||
|
|
|
||
|
|
from collections.abc import AsyncIterator
|
||
|
|
from typing import Any
|
||
|
|
|
||
|
|
import pytest
|
||
|
|
from fastapi import Request
|
||
|
|
from fastapi.testclient import TestClient
|
||
|
|
|
||
|
|
from app.api.dependencies.auth import build_request_context
|
||
|
|
from app.api.dependencies.database import get_session
|
||
|
|
from app.core.config import get_settings
|
||
|
|
from app.core.contracts import RequestContext
|
||
|
|
from app.main import create_app
|
||
|
|
|
||
|
|
MESSAGES = "/api/v1/conversations/session-1/messages"
|
||
|
|
RUN_DETAIL = "/api/v1/agent-runs/run-1"
|
||
|
|
|
||
|
|
|
||
|
|
class FakeBackend:
|
||
|
|
"""替身计数后端:固定返回给定的 `(计数, 剩余秒数)`,或 `None`(后端不可用)。"""
|
||
|
|
|
||
|
|
def __init__(self, result: tuple[int, int] | None) -> None:
|
||
|
|
self.result = result
|
||
|
|
self.calls: list[tuple[str, int]] = []
|
||
|
|
|
||
|
|
async def increment(self, key: str, window_seconds: int) -> tuple[int, int] | None:
|
||
|
|
self.calls.append((key, window_seconds))
|
||
|
|
return self.result
|
||
|
|
|
||
|
|
|
||
|
|
class RecordingService:
|
||
|
|
def __init__(self, session: Any) -> None:
|
||
|
|
del session
|
||
|
|
|
||
|
|
async def messages(
|
||
|
|
self, session_id: str, context: RequestContext, limit: int, before: int | None = None
|
||
|
|
) -> dict[str, object]:
|
||
|
|
return {"data": []}
|
||
|
|
|
||
|
|
|
||
|
|
async def override_session() -> AsyncIterator[Any]:
|
||
|
|
yield object()
|
||
|
|
|
||
|
|
|
||
|
|
def build_client(
|
||
|
|
monkeypatch: pytest.MonkeyPatch,
|
||
|
|
*,
|
||
|
|
backend: FakeBackend | None,
|
||
|
|
user_id: str = "9001",
|
||
|
|
max_requests: int = 2,
|
||
|
|
enabled: bool = True,
|
||
|
|
) -> TestClient:
|
||
|
|
monkeypatch.setattr(
|
||
|
|
"app.api.dependencies.rate_limit.get_settings",
|
||
|
|
lambda: get_settings().model_copy(update={
|
||
|
|
"rate_limit_enabled": enabled,
|
||
|
|
"rate_limit_window_seconds": 60,
|
||
|
|
"rate_limit_max_requests": max_requests,
|
||
|
|
"rate_limit_key_prefix": "test:rate_limit",
|
||
|
|
}),
|
||
|
|
)
|
||
|
|
if backend is not None:
|
||
|
|
monkeypatch.setattr(
|
||
|
|
"app.api.dependencies.rate_limit.get_counter_backend", lambda: backend
|
||
|
|
)
|
||
|
|
monkeypatch.setattr("app.api.controllers.conversations.ConversationService", RecordingService)
|
||
|
|
|
||
|
|
application = create_app()
|
||
|
|
|
||
|
|
async def context(request: Request) -> RequestContext:
|
||
|
|
built = RequestContext(user_id=user_id, trace_id="trace-rate",
|
||
|
|
permissions=("conversation:read", "agent:run"))
|
||
|
|
request.state.request_context = built
|
||
|
|
return built
|
||
|
|
|
||
|
|
application.dependency_overrides[build_request_context] = context
|
||
|
|
application.dependency_overrides[get_session] = override_session
|
||
|
|
return TestClient(application)
|
||
|
|
|
||
|
|
|
||
|
|
def test_request_within_limit_is_allowed(monkeypatch: pytest.MonkeyPatch) -> None:
|
||
|
|
backend = FakeBackend((2, 30)) # 恰好等于阈值:文档语义是"超过"才拒绝
|
||
|
|
expected_key = "test:rate_limit:9001:GET:/api/v1/conversations/{session_id}/messages"
|
||
|
|
|
||
|
|
with build_client(monkeypatch, backend=backend, max_requests=2) as client:
|
||
|
|
response = client.get(MESSAGES)
|
||
|
|
|
||
|
|
assert response.status_code == 200
|
||
|
|
assert backend.calls == [(expected_key, 60)]
|
||
|
|
|
||
|
|
|
||
|
|
def test_request_over_limit_returns_429_with_code_and_retry_after(
|
||
|
|
monkeypatch: pytest.MonkeyPatch
|
||
|
|
) -> None:
|
||
|
|
backend = FakeBackend((3, 17))
|
||
|
|
|
||
|
|
with build_client(monkeypatch, backend=backend, max_requests=2) as client:
|
||
|
|
response = client.get(MESSAGES)
|
||
|
|
|
||
|
|
assert response.status_code == 429
|
||
|
|
assert response.headers["Retry-After"] == "17"
|
||
|
|
body = response.json()
|
||
|
|
assert set(body) == {"error", "meta"}
|
||
|
|
assert body["error"]["code"] == "RATE_LIMITED"
|
||
|
|
assert body["error"]["retryable"] is True # 文档 §3.6:该码可重试
|
||
|
|
assert body["error"]["field_errors"] == []
|
||
|
|
assert body["meta"]["trace_id"] == "trace-rate"
|
||
|
|
|
||
|
|
|
||
|
|
def test_redis_unavailable_degrades_to_allow(monkeypatch: pytest.MonkeyPatch) -> None:
|
||
|
|
"""Redis 不可用时必须放行:保护措施不能变成全量拒绝。"""
|
||
|
|
backend = FakeBackend(None)
|
||
|
|
|
||
|
|
with build_client(monkeypatch, backend=backend) as client:
|
||
|
|
response = client.get(MESSAGES)
|
||
|
|
|
||
|
|
assert response.status_code == 200
|
||
|
|
assert len(backend.calls) == 1 # 确实尝试过判定,只是判定不出来
|
||
|
|
|
||
|
|
|
||
|
|
def test_disabled_limit_does_not_touch_backend(monkeypatch: pytest.MonkeyPatch) -> None:
|
||
|
|
backend = FakeBackend((99, 5))
|
||
|
|
|
||
|
|
with build_client(monkeypatch, backend=backend, enabled=False) as client:
|
||
|
|
response = client.get(MESSAGES)
|
||
|
|
|
||
|
|
assert response.status_code == 200
|
||
|
|
assert backend.calls == []
|
||
|
|
|
||
|
|
|
||
|
|
def test_authentication_precedes_rate_limiting(monkeypatch: pytest.MonkeyPatch) -> None:
|
||
|
|
"""未带令牌 + 已超限:必须仍是 401(限流不得掩盖鉴权失败)。"""
|
||
|
|
application = create_app()
|
||
|
|
monkeypatch.setattr(
|
||
|
|
"app.api.dependencies.rate_limit.get_settings",
|
||
|
|
lambda: get_settings().model_copy(update={"rate_limit_max_requests": 1}),
|
||
|
|
)
|
||
|
|
monkeypatch.setattr(
|
||
|
|
"app.api.dependencies.rate_limit.get_counter_backend", lambda: FakeBackend((99, 5))
|
||
|
|
)
|
||
|
|
|
||
|
|
with TestClient(application) as client:
|
||
|
|
response = client.get(MESSAGES)
|
||
|
|
|
||
|
|
assert response.status_code == 401
|
||
|
|
assert response.json()["error"]["code"] == "AUTHENTICATION_REQUIRED"
|
||
|
|
|
||
|
|
|
||
|
|
def test_counter_key_uses_route_template_and_user(monkeypatch: pytest.MonkeyPatch) -> None:
|
||
|
|
"""计数维度:同一路由模板下不同 `run_id` 共用一个计数器,不同用户互相隔离。"""
|
||
|
|
first = FakeBackend((1, 60))
|
||
|
|
second = FakeBackend((1, 60))
|
||
|
|
|
||
|
|
with build_client(monkeypatch, backend=first, user_id="9001") as client:
|
||
|
|
client.get("/api/v1/agent-runs/run-1")
|
||
|
|
client.get("/api/v1/agent-runs/run-2")
|
||
|
|
with build_client(monkeypatch, backend=second, user_id="9002") as client:
|
||
|
|
client.get(RUN_DETAIL)
|
||
|
|
|
||
|
|
assert first.calls == [
|
||
|
|
("test:rate_limit:9001:GET:/api/v1/agent-runs/{run_id}", 60),
|
||
|
|
("test:rate_limit:9001:GET:/api/v1/agent-runs/{run_id}", 60),
|
||
|
|
]
|
||
|
|
assert second.calls == [("test:rate_limit:9002:GET:/api/v1/agent-runs/{run_id}", 60)]
|