from pathlib import Path from types import SimpleNamespace from unittest.mock import patch from fastapi.testclient import TestClient from database import SessionLocal from main import app from models import Avatar, KnowledgeChunk, KnowledgeDoc, QAPair from routers.knowledge import _doc_payload client = TestClient(app) def test_doc_payload_reports_whether_the_persisted_file_exists(tmp_path: Path): avatar_id = "avatar-1" stored_name = "knowledge.md" doc = SimpleNamespace( avatar_id=avatar_id, file_url=f"/api/files/{avatar_id}/{stored_name}", to_dict=lambda: {"id": "doc-1", "fileUrl": f"/api/files/{avatar_id}/{stored_name}"}, ) stored_dir = tmp_path / avatar_id stored_dir.mkdir() stored_file = stored_dir / stored_name with patch("routers.knowledge.UPLOAD_DIR", str(tmp_path)): assert _doc_payload(doc)["filePresent"] is False stored_file.write_text("knowledge", encoding="utf-8") assert _doc_payload(doc)["filePresent"] is True def test_upload_marks_vectorization_failure_instead_of_staying_processing( tmp_path: Path, authorization_context, ): context = authorization_context with ( patch("routers.knowledge.UPLOAD_DIR", str(tmp_path)), patch("routers.knowledge.embeddings.embed", side_effect=RuntimeError("provider unavailable")), ): response = client.post( f"/api/avatar/{context['avatar'].id}/knowledge/docs", headers=context["owner_headers"], files={"file": ("knowledge.md", b"# Knowledge\n\nTest content", "text/markdown")}, ) payload = response.json()["data"] assert payload["status"] == "failed" assert payload["vectorized"] is False assert payload["chunkCount"] == 0 db = SessionLocal() try: stored = db.query(KnowledgeDoc).filter(KnowledgeDoc.id == payload["id"]).one() assert stored.status == "failed" assert db.query(KnowledgeChunk).filter(KnowledgeChunk.doc_id == stored.id).count() == 0 db.delete(stored) db.commit() finally: db.close() def test_markdown_upload_commits_ready_document_and_chunks_together( tmp_path: Path, authorization_context, ): context = authorization_context with ( patch("routers.knowledge.UPLOAD_DIR", str(tmp_path)), patch("routers.knowledge.embeddings.embed", return_value=[[1.0, 0.0]]), ): response = client.post( f"/api/avatar/{context['avatar'].id}/knowledge/docs", headers=context["owner_headers"], files={"file": ("knowledge.md", b"# Knowledge\n\nTest content", "text/markdown")}, ) payload = response.json()["data"] assert payload["status"] == "ready" assert payload["vectorized"] is True assert payload["chunkCount"] == 1 db = SessionLocal() try: stored = db.query(KnowledgeDoc).filter(KnowledgeDoc.id == payload["id"]).one() assert stored.status == "ready" assert db.query(KnowledgeChunk).filter(KnowledgeChunk.doc_id == stored.id).count() == 1 db.query(KnowledgeChunk).filter(KnowledgeChunk.doc_id == stored.id).delete() db.delete(stored) db.commit() finally: db.close() def test_each_avatar_has_an_independent_document_and_qa_scope(authorization_context): context = authorization_context first_avatar_id = context["avatar"].id second_avatar_id = f"knowledge-second-{context['suffix']}" first_doc_id = f"knowledge-first-doc-{context['suffix']}" second_doc_id = f"knowledge-second-doc-{context['suffix']}" first_qa_id = f"knowledge-first-qa-{context['suffix']}" second_qa_id = f"knowledge-second-qa-{context['suffix']}" db = SessionLocal() try: db.add_all( [ Avatar( id=second_avatar_id, owner_id=context["owner"].huihui_user_id, name="独立知识库分身", status="active", config={}, ), KnowledgeDoc( id=first_doc_id, avatar_id=first_avatar_id, filename="first.md", status="ready", vectorized=True, ), KnowledgeDoc( id=second_doc_id, avatar_id=second_avatar_id, filename="second.md", status="ready", vectorized=True, ), QAPair( id=first_qa_id, avatar_id=first_avatar_id, question="第一个分身问题", answer="第一个分身答案", ), QAPair( id=second_qa_id, avatar_id=second_avatar_id, question="第二个分身问题", answer="第二个分身答案", ), ] ) db.commit() finally: db.close() try: first_docs = client.get( f"/api/avatar/{first_avatar_id}/knowledge/docs", headers=context["owner_headers"], ).json()["data"] second_docs = client.get( f"/api/avatar/{second_avatar_id}/knowledge/docs", headers=context["owner_headers"], ).json()["data"] first_qa = client.get( f"/api/avatar/{first_avatar_id}/knowledge/qa", headers=context["owner_headers"], ).json()["data"] second_qa = client.get( f"/api/avatar/{second_avatar_id}/knowledge/qa", headers=context["owner_headers"], ).json()["data"] assert [item["id"] for item in first_docs if item["id"] == first_doc_id] == [first_doc_id] assert second_doc_id not in {item["id"] for item in first_docs} assert [item["id"] for item in second_docs] == [second_doc_id] assert first_qa_id in {item["id"] for item in first_qa} assert second_qa_id not in {item["id"] for item in first_qa} assert [item["id"] for item in second_qa] == [second_qa_id] finally: db = SessionLocal() try: db.query(QAPair).filter(QAPair.id.in_([first_qa_id, second_qa_id])).delete( synchronize_session=False ) db.query(KnowledgeDoc).filter( KnowledgeDoc.id.in_([first_doc_id, second_doc_id]) ).delete(synchronize_session=False) db.query(Avatar).filter(Avatar.id == second_avatar_id).delete() db.commit() finally: db.close()