95 lines
3.2 KiB
Python
95 lines
3.2 KiB
Python
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 KnowledgeChunk, KnowledgeDoc
|
|
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()
|