PlantAI-rag-new/sdk/app/test_sessions_api.py

"""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, "capabilities": {}}

    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, "capabilities": {}}

    # 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