244 lines
8.5 KiB
Python
244 lines
8.5 KiB
Python
|
"""app/test_sessions_api.py —— 会话/消息路由集成单测（API.md §5 / §1 信封）。
|
|||
|
|
|
|||
|
|
TestClient + 临时 SQLite（dependency_overrides 指向临时 Database），
|
|||
|
|
不触发 app lifespan（不跑 init_db/上游探测）。
|
|||
|
|
"""
|
|||
|
|
from __future__ import annotations
|
|||
|
|
|
|||
|
|
import pytest
|
|||
|
|
from fastapi.testclient import TestClient
|
|||
|
|
|
|||
|
|
from app.db import Database, dumps_json, get_db
|
|||
|
|
from app.main import app
|
|||
|
|
from app.upstream import envelope
|
|||
|
|
|
|||
|
|
|
|||
|
|
@pytest.fixture()
|
|||
|
|
async def db(tmp_path):
|
|||
|
|
database = Database(tmp_path / "api" / "plantai.db")
|
|||
|
|
await database.connect()
|
|||
|
|
yield database
|
|||
|
|
await database.close()
|
|||
|
|
|
|||
|
|
|
|||
|
|
@pytest.fixture()
|
|||
|
|
def client(db):
|
|||
|
|
app.dependency_overrides[get_db] = lambda: db
|
|||
|
|
try:
|
|||
|
|
yield TestClient(app)
|
|||
|
|
finally:
|
|||
|
|
app.dependency_overrides.clear()
|
|||
|
|
|
|||
|
|
|
|||
|
|
# ---------------------------------------------------------------------------
|
|||
|
|
# 错误信封构造器（§1）
|
|||
|
|
# ---------------------------------------------------------------------------
|
|||
|
|
|
|||
|
|
@pytest.mark.parametrize(
|
|||
|
|
"code,expected_status",
|
|||
|
|
[
|
|||
|
|
("UPSTREAM_DOWN", 503),
|
|||
|
|
("UPSTREAM_ERROR", 502),
|
|||
|
|
("TIMEOUT", 504),
|
|||
|
|
("INVALID_CITATION", 502),
|
|||
|
|
("NOT_FOUND", 404),
|
|||
|
|
("VALIDATION", 422),
|
|||
|
|
("INTERNAL", 500),
|
|||
|
|
],
|
|||
|
|
)
|
|||
|
|
def test_envelope_defaults(code, expected_status):
|
|||
|
|
body, status = envelope(code)
|
|||
|
|
assert status == expected_status
|
|||
|
|
assert body["ok"] is False
|
|||
|
|
assert body["error"]["code"] == code
|
|||
|
|
assert isinstance(body["error"]["message"], str) and body["error"]["message"]
|
|||
|
|
assert "degraded" in body["error"]
|
|||
|
|
|
|||
|
|
|
|||
|
|
def test_envelope_custom():
|
|||
|
|
body, status = envelope("UPSTREAM_DOWN", message="自定义", degraded=True, status=503)
|
|||
|
|
assert body["error"]["message"] == "自定义"
|
|||
|
|
assert body["error"]["degraded"] is True
|
|||
|
|
|
|||
|
|
|
|||
|
|
# ---------------------------------------------------------------------------
|
|||
|
|
# 会话 CRUD 全链路
|
|||
|
|
# ---------------------------------------------------------------------------
|
|||
|
|
|
|||
|
|
def test_session_crud_full_flow(client):
|
|||
|
|
# POST（默认标题 + 自定义标题）
|
|||
|
|
r = client.post("/api/sessions", json={})
|
|||
|
|
assert r.status_code == 201, r.text
|
|||
|
|
sess = r.json()
|
|||
|
|
assert sess["title"] == "新会话"
|
|||
|
|
assert sess["status"] == "idle"
|
|||
|
|
assert sess["message_count"] == 0
|
|||
|
|
assert sess["meta"] == {"title_manual": False}
|
|||
|
|
|
|||
|
|
r2 = client.post("/api/sessions", json={"title": " 水稻SPL "})
|
|||
|
|
sess2 = r2.json()
|
|||
|
|
assert sess2["title"] == "水稻SPL"
|
|||
|
|
|
|||
|
|
# GET 列表（updated_at 倒序，最新在前）
|
|||
|
|
r = client.get("/api/sessions")
|
|||
|
|
assert r.status_code == 200
|
|||
|
|
items = r.json()
|
|||
|
|
assert [s["id"] for s in items][:2] == [sess2["id"], sess["id"]]
|
|||
|
|
|
|||
|
|
# POST user 消息
|
|||
|
|
r = client.post(f"/api/sessions/{sess['id']}/messages", json={"role": "user", "text": "问题原文"})
|
|||
|
|
assert r.status_code == 201, r.text
|
|||
|
|
msg = r.json()
|
|||
|
|
assert msg["role"] == "user" and msg["status"] == "done"
|
|||
|
|
assert msg["text"] == "问题原文"
|
|||
|
|
assert msg["blocks"] == [] and msg["raw"] is None
|
|||
|
|
|
|||
|
|
# POST assistant 流式骨架
|
|||
|
|
r = client.post(
|
|||
|
|
f"/api/sessions/{sess['id']}/messages",
|
|||
|
|
json={
|
|||
|
|
"role": "assistant", "kind": "qa", "status": "streaming",
|
|||
|
|
"blocks": [{"type": "pipeline", "status": "running", "steps": []}],
|
|||
|
|
},
|
|||
|
|
)
|
|||
|
|
assert r.status_code == 201, r.text
|
|||
|
|
skeleton = r.json()
|
|||
|
|
assert skeleton["status"] == "streaming"
|
|||
|
|
assert skeleton["blocks"][0]["type"] == "pipeline"
|
|||
|
|
assert "status" in skeleton["blocks"][0]
|
|||
|
|
|
|||
|
|
# 消息活动后会话状态 running、message_count 更新
|
|||
|
|
r = client.get(f"/api/sessions/{sess['id']}")
|
|||
|
|
detail = r.json()
|
|||
|
|
assert detail["status"] == "running"
|
|||
|
|
assert detail["message_count"] == 2
|
|||
|
|
assert [m["role"] for m in detail["messages"]] == ["user", "assistant"]
|
|||
|
|
|
|||
|
|
# PATCH 消息（流式检查点：status/blocks/raw）
|
|||
|
|
r = client.patch(
|
|||
|
|
f"/api/sessions/{sess['id']}/messages/{skeleton['id']}",
|
|||
|
|
json={
|
|||
|
|
"status": "done",
|
|||
|
|
"blocks": [{"type": "text", "html": "<p>最终答案</p>"}],
|
|||
|
|
"raw": {"answer": "最终答案"},
|
|||
|
|
},
|
|||
|
|
)
|
|||
|
|
assert r.status_code == 200, r.text
|
|||
|
|
updated = r.json()
|
|||
|
|
assert updated["status"] == "done"
|
|||
|
|
assert updated["blocks"] == [{"type": "text", "html": "<p>最终答案</p>"}]
|
|||
|
|
assert updated["raw"] == {"answer": "最终答案"}
|
|||
|
|
|
|||
|
|
# PATCH 会话标题 → title_manual=true
|
|||
|
|
r = client.patch(f"/api/sessions/{sess['id']}", json={"title": "手工标题"})
|
|||
|
|
assert r.status_code == 200
|
|||
|
|
assert r.json()["meta"] == {"title_manual": True}
|
|||
|
|
|
|||
|
|
# PATCH 会话状态
|
|||
|
|
r = client.patch(f"/api/sessions/{sess['id']}", json={"status": "done"})
|
|||
|
|
assert r.status_code == 200
|
|||
|
|
assert r.json()["status"] == "done"
|
|||
|
|
|
|||
|
|
# DELETE 204 + 级联
|
|||
|
|
r = client.delete(f"/api/sessions/{sess['id']}")
|
|||
|
|
assert r.status_code == 204
|
|||
|
|
r = client.get(f"/api/sessions/{sess['id']}")
|
|||
|
|
assert r.status_code == 404
|
|||
|
|
body = r.json()
|
|||
|
|
assert body["ok"] is False
|
|||
|
|
assert body["error"]["code"] == "NOT_FOUND"
|
|||
|
|
|
|||
|
|
|
|||
|
|
def test_not_found_envelopes(client):
|
|||
|
|
cases = [
|
|||
|
|
("get", "/api/sessions/nope", None),
|
|||
|
|
("patch", "/api/sessions/nope", {"title": "x"}),
|
|||
|
|
("delete", "/api/sessions/nope", None),
|
|||
|
|
("post", "/api/sessions/nope/messages", {"role": "user", "text": "hi"}),
|
|||
|
|
("get", "/api/sessions/nope/sources", None),
|
|||
|
|
]
|
|||
|
|
for method, url, body in cases:
|
|||
|
|
kwargs = {"json": body} if body is not None else {}
|
|||
|
|
r = getattr(client, method)(url, **kwargs)
|
|||
|
|
assert r.status_code == 404, (method, url)
|
|||
|
|
payload = r.json()
|
|||
|
|
assert payload["ok"] is False
|
|||
|
|
assert payload["error"]["code"] == "NOT_FOUND"
|
|||
|
|
|
|||
|
|
r = client.patch("/api/sessions/nope/messages/nope-msg", json={"status": "done"})
|
|||
|
|
assert r.status_code == 404
|
|||
|
|
assert r.json()["error"]["code"] == "NOT_FOUND"
|
|||
|
|
|
|||
|
|
|
|||
|
|
def test_message_404_when_session_exists_but_message_missing(client):
|
|||
|
|
r = client.post("/api/sessions", json={"title": "t"})
|
|||
|
|
sid = r.json()["id"]
|
|||
|
|
r = client.patch(f"/api/sessions/{sid}/messages/ghost", json={"status": "done"})
|
|||
|
|
assert r.status_code == 404
|
|||
|
|
assert r.json()["error"]["code"] == "NOT_FOUND"
|
|||
|
|
|
|||
|
|
|
|||
|
|
def test_empty_title_falls_back_to_default(client):
|
|||
|
|
r = client.post("/api/sessions", json={"title": " "})
|
|||
|
|
assert r.json()["title"] == "新会话"
|
|||
|
|
|
|||
|
|
|
|||
|
|
# ---------------------------------------------------------------------------
|
|||
|
|
# 来源聚合（GET /api/sessions/{id}/sources）
|
|||
|
|
# ---------------------------------------------------------------------------
|
|||
|
|
|
|||
|
|
def _source_item(i: str, paper: str | None, title: str) -> dict:
|
|||
|
|
return {
|
|||
|
|
"id": i, "title": title, "paper_id": paper, "chunk_id": None,
|
|||
|
|
"pmid": None, "pmcid": None, "doi": None, "url": None,
|
|||
|
|
"snippet": "…", "score": 0.9, "locator": None, "section": None,
|
|||
|
|
"support_count": None, "type": "assertion", "metadata": {},
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
|
|||
|
|
async def test_sources_aggregation(client, db):
|
|||
|
|
r = client.post("/api/sessions", json={"title": "聚合"})
|
|||
|
|
sid = r.json()["id"]
|
|||
|
|
|
|||
|
|
# 两个助手消息各带一个 sources 块
|
|||
|
|
blocks1 = [
|
|||
|
|
{
|
|||
|
|
"type": "sources",
|
|||
|
|
"items": [
|
|||
|
|
_source_item("a1", "PAPER_A", "Paper A"),
|
|||
|
|
_source_item("b1", "PAPER_B", "Paper B"),
|
|||
|
|
_source_item("x1", None, "No paper id"),
|
|||
|
|
],
|
|||
|
|
}
|
|||
|
|
]
|
|||
|
|
blocks2 = [
|
|||
|
|
{
|
|||
|
|
"type": "sources",
|
|||
|
|
"items": [
|
|||
|
|
_source_item("a2", "PAPER_A", "Paper A"),
|
|||
|
|
_source_item("x2", None, "No paper id"),
|
|||
|
|
],
|
|||
|
|
}
|
|||
|
|
]
|
|||
|
|
await db.execute(
|
|||
|
|
"INSERT INTO messages (id, session_id, role, kind, status, blocks, raw, text, "
|
|||
|
|
"created_at, updated_at) VALUES "
|
|||
|
|
"('m1', ?, 'assistant', 'qa', 'done', ?, NULL, NULL, '2026-01-01T00:00:00Z', '2026-01-01T00:00:00Z'), "
|
|||
|
|
"('m2', ?, 'assistant', 'qa', 'done', ?, NULL, NULL, '2026-01-01T00:00:00Z', '2026-01-01T00:00:00Z')",
|
|||
|
|
(sid, dumps_json(blocks1), sid, dumps_json(blocks2)),
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
r = client.get(f"/api/sessions/{sid}/sources")
|
|||
|
|
assert r.status_code == 200, r.text
|
|||
|
|
body = r.json()
|
|||
|
|
assert body["session_id"] == sid
|
|||
|
|
# PAPER_A 出现 2 次（2 个消息各 1 条），PAPER_B 1 次；无 paper_id 的按 id 各 1 次
|
|||
|
|
assert body["total"] == 5
|
|||
|
|
by_key = {(s["paper_id"] or s["id"]): s for s in body["sources"]}
|
|||
|
|
assert by_key["PAPER_A"]["count"] == 2
|
|||
|
|
assert by_key["PAPER_A"]["title"] == "Paper A"
|
|||
|
|
assert by_key["PAPER_B"]["count"] == 1
|
|||
|
|
assert by_key["x1"]["count"] == 1
|
|||
|
|
assert by_key["x2"]["count"] == 1
|
|||
|
|
assert len(body["sources"]) == 4
|