91 lines
3.2 KiB
Python
91 lines
3.2 KiB
Python
# requires: fastapi, httpx
|
|
import json
|
|
from collections.abc import AsyncIterator
|
|
|
|
from fastapi.testclient import TestClient
|
|
|
|
from app import Message, app, get_llm
|
|
|
|
AUTH = {"X-API-Key": "demo-key"}
|
|
BODY = {"messages": [{"role": "user", "content": "你好"}]}
|
|
|
|
|
|
class ScriptedLLM:
|
|
"""测试替身：按脚本返回 token，可以模拟中途失败。"""
|
|
|
|
def __init__(self, tokens: list[str], fail_at: int | None = None) -> None:
|
|
self.tokens, self.fail_at = tokens, fail_at
|
|
self.calls: list[list[Message]] = []
|
|
|
|
async def stream(self, messages: list[Message], temperature: float) -> AsyncIterator[str]:
|
|
self.calls.append(messages)
|
|
for i, t in enumerate(self.tokens):
|
|
if i == self.fail_at:
|
|
raise RuntimeError("上游模型 500")
|
|
yield t
|
|
|
|
|
|
def client_with(llm) -> TestClient:
|
|
app.dependency_overrides[get_llm] = lambda: llm
|
|
return TestClient(app, raise_server_exceptions=False)
|
|
|
|
|
|
def parse_sse(text: str) -> list[tuple[str, str]]:
|
|
events = []
|
|
for block in text.strip().split("\n\n"):
|
|
name, data = "message", ""
|
|
for line in block.splitlines():
|
|
if line.startswith("event: "):
|
|
name = line[7:]
|
|
elif line.startswith("data: "):
|
|
data = line[6:]
|
|
events.append((name, data))
|
|
return events
|
|
|
|
|
|
def teardown_function():
|
|
app.dependency_overrides.clear()
|
|
|
|
|
|
def test_auth_required():
|
|
r = client_with(ScriptedLLM(["x"])).post("/v1/chat", json=BODY)
|
|
assert r.status_code == 401 and r.json()["code"] == "HTTP_401" and r.json()["request_id"]
|
|
|
|
|
|
def test_validation_error_format():
|
|
r = client_with(ScriptedLLM(["x"])).post("/v1/chat", headers=AUTH,
|
|
json={"messages": [{"role": "robot", "content": ""}], "temperature": 5})
|
|
assert r.status_code == 422
|
|
fields = {d["field"] for d in r.json()["details"]}
|
|
assert {"messages.0.role", "messages.0.content", "temperature"} <= fields
|
|
|
|
|
|
def test_non_stream():
|
|
llm = ScriptedLLM(["你", "好"])
|
|
r = client_with(llm).post("/v1/chat", headers={**AUTH, "X-Request-ID": "rid-1"}, json={**BODY, "stream": False})
|
|
assert r.status_code == 200
|
|
assert r.json() == {"request_id": "rid-1", "content": "你好"}
|
|
assert r.headers["X-Request-ID"] == "rid-1"
|
|
assert llm.calls[0][0].content == "你好"
|
|
|
|
|
|
def test_stream_sse():
|
|
r = client_with(ScriptedLLM(["A", "B", "C"])).post("/v1/chat", headers=AUTH, json=BODY)
|
|
assert r.status_code == 200 and r.headers["content-type"].startswith("text/event-stream")
|
|
events = parse_sse(r.text)
|
|
assert [json.loads(d)["delta"] for n, d in events[:-1]] == ["A", "B", "C"]
|
|
assert events[-1] == ("message", "[DONE]")
|
|
|
|
|
|
def test_stream_error_event():
|
|
r = client_with(ScriptedLLM(["A", "B"], fail_at=1)).post("/v1/chat", headers=AUTH, json=BODY)
|
|
events = parse_sse(r.text)
|
|
assert events[0][0] == "message" and events[-1][0] == "error"
|
|
assert json.loads(events[-1][1])["message"] == "生成失败"
|
|
|
|
|
|
def test_internal_error_hides_details():
|
|
r = client_with(ScriptedLLM([])).get("/boom")
|
|
assert r.status_code == 500
|
|
assert r.json()["message"] == "服务器内部错误" and "postgres" not in r.text
|