huihuiSquare/digital-avatar-app/backend/tests/test_knowledge_storage.py

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()