import json import os from pathlib import Path import socket import sqlite3 import subprocess import time import httpx import pytest from pymilvus import MilvusClient pytestmark = pytest.mark.live ROOT = Path(__file__).resolve().parents[3] FIXTURES = ROOT / "backend" / "tests" / "fixtures" FIXTURE_PATHS = tuple( FIXTURES / name for name in ("knowledge.md", "knowledge.txt", "knowledge.pdf") ) def _port() -> int: with socket.socket() as sock: sock.bind(("127.0.0.1", 0)) return int(sock.getsockname()[1]) def _start_server(port: int, env: dict[str, str]) -> subprocess.Popen[str]: process = subprocess.Popen( [ str(ROOT / "backend" / ".venv" / "bin" / "python"), "-m", "uvicorn", "kbqa.main:app", "--app-dir", "backend/src", "--host", "127.0.0.1", "--port", str(port), ], cwd=ROOT, env=env, stdout=subprocess.DEVNULL, stderr=subprocess.STDOUT, text=True, ) base_url = f"http://127.0.0.1:{port}" for _ in range(120): if process.poll() is not None: raise RuntimeError("live Uvicorn exited during startup") try: if httpx.get(f"{base_url}/health", timeout=1).status_code == 200: return process except httpx.HTTPError: pass time.sleep(0.25) process.terminate() raise TimeoutError("live Uvicorn did not become healthy") def _stop_server(process: subprocess.Popen[str]) -> None: process.terminate() try: process.wait(timeout=15) except subprocess.TimeoutExpired: process.kill() process.wait(timeout=5) def _stream(client: httpx.Client, session_id: str, query: str) -> list[tuple[str, dict]]: events: list[tuple[str, dict]] = [] with client.stream( "POST", "/api/v1/chat/stream", json={"session_id": session_id, "query": query}, timeout=180 ) as response: response.raise_for_status() event_name = "" for line in response.iter_lines(): if line.startswith("event: "): event_name = line.removeprefix("event: ") elif line.startswith("data: "): events.append((event_name, json.loads(line.removeprefix("data: ")))) return events def _answer(events: list[tuple[str, dict]]) -> str: return "".join(data["content"] for name, data in events if name == "token") def _assert_flow(events: list[tuple[str, dict]]) -> None: stages = [ f"status:{data['phase']}" if name == "status" else name for name, data in events if name != "ping" ] expected = ["status:researching", "sources", "status:answering", "token", "done"] positions = [stages.index(stage) for stage in expected] assert positions == sorted(positions), stages def test_real_model_pipeline_survives_restart_and_cleans_up(tmp_path) -> None: if os.environ.get("KBQA_RUN_LIVE_TESTS") != "1": pytest.skip("set KBQA_RUN_LIVE_TESTS=1 to call real models") database = tmp_path / "kbqa.sqlite3" milvus = tmp_path / "milvus.db" raw = tmp_path / "raw" env = os.environ.copy() env.update( { "KBQA_DATABASE_URL": f"sqlite+aiosqlite:///{database}", "KBQA_MILVUS_URI": str(milvus), "KBQA_RAW_DATA_DIR": str(raw), "KBQA_RUN_LIVE_TESTS": "1", } ) subprocess.run( [ str(ROOT / "backend" / ".venv" / "bin" / "alembic"), "-c", "backend/alembic.ini", "upgrade", "head", ], cwd=ROOT, env=env, check=True, capture_output=True, text=True, ) port = _port() process = _start_server(port, env) base_url = f"http://127.0.0.1:{port}" document_ids: list[str] = [] try: with httpx.Client(base_url=base_url, timeout=30) as client: for path in FIXTURE_PATHS: with path.open("rb") as file_handle: response = client.post( "/api/v1/documents", files={"file": (path.name, file_handle)} ) response.raise_for_status() document_ids.append(response.json()["id"]) deadline = time.monotonic() + 120 while time.monotonic() < deadline: documents = client.get("/api/v1/documents", params={"limit": 100}).json()["items"] if all(item["status"] == "indexed" for item in documents): break assert not any(item["status"] == "failed" for item in documents), documents time.sleep(0.5) else: pytest.fail("documents did not finish indexing") assert len(documents) == 3 assert all(item["chunk_count"] > 0 for item in documents) session = client.post("/api/v1/sessions", json={"title": "真实模型测试"}).json() session_id = session["id"] first = _stream(client, session_id, "星槐项目的内部代号和负责人分别是什么？") _assert_flow(first) first_answer = _answer(first) assert "ORCHID-731" in first_answer assert "林岚" in first_answer first_sources = [ source for name, data in first if name == "sources" for source in data["sources"] ] assert first_sources assert {source["filename"] for source in first_sources} <= { path.name for path in FIXTURE_PATHS } chunks_by_document: dict[str, set[str]] = {} for source in first_sources: document_id = source["document_id"] if document_id not in chunks_by_document: page = client.get( f"/api/v1/documents/{document_id}/chunks", params={"limit": 100} ).json() chunks_by_document[document_id] = {item["id"] for item in page["items"]} assert source["chunk_id"] in chunks_by_document[document_id] follow_up = _stream(client, session_id, "她负责的项目首次评审日期是什么？") _assert_flow(follow_up) follow_up_answer = _answer(follow_up) assert "2026" in follow_up_answer and "18" in follow_up_answer absent = _stream(client, session_id, "星槐项目部署在哪个月球基地？") _assert_flow(absent) assert any(marker in _answer(absent) for marker in ("不足", "没有", "未检索")) history = client.get( f"/api/v1/sessions/{session_id}/messages", params={"limit": 100} ).json() assert history["total"] == 6 assert all(item["status"] == "completed" for item in history["items"]) assert history["items"][1]["sources"] finally: _stop_server(process) process = _start_server(port, env) try: with httpx.Client(base_url=base_url, timeout=30) as client: assert client.get("/api/v1/documents").json()["total"] == 3 history = client.get( f"/api/v1/sessions/{session_id}/messages", params={"limit": 100} ).json() assert history["total"] == 6 after_restart = _stream(client, session_id, "青屿实验的校准周期是多久？") _assert_flow(after_restart) assert "14" in _answer(after_restart) for document_id in document_ids: assert client.delete(f"/api/v1/documents/{document_id}").status_code == 204 assert client.delete(f"/api/v1/sessions/{session_id}").status_code == 204 assert client.get("/api/v1/documents").json()["total"] == 0 assert client.get(f"/api/v1/sessions/{session_id}/messages").status_code == 404 finally: _stop_server(process) with sqlite3.connect(database) as connection: for table in ( "documents", "index_jobs", "chunks", "chat_sessions", "messages", "message_sources", ): assert connection.execute(f"SELECT count(*) FROM {table}").fetchone()[0] == 0 vector_client = MilvusClient(str(milvus)) try: vector_client.load_collection("kbqa_chunks") assert int(vector_client.get_collection_stats("kbqa_chunks")["row_count"]) == 0 finally: vector_client.close() assert not list(raw.glob("*"))