148 lines
4.9 KiB
Python
148 lines
4.9 KiB
Python
import tempfile
|
|
import unittest
|
|
from unittest.mock import AsyncMock
|
|
|
|
from api.chat.knowledge import (
|
|
confirm_document_upload,
|
|
delete_document,
|
|
get_document,
|
|
list_documents,
|
|
preview_document_upload,
|
|
)
|
|
from service.knowledge_base.upload import KnowledgeUploadService
|
|
from types import SimpleNamespace
|
|
|
|
|
|
class FakeUpload:
|
|
def __init__(self, filename, content):
|
|
self.filename = filename
|
|
self._content = content
|
|
|
|
async def read(self):
|
|
return self._content
|
|
|
|
|
|
class FakeRequest:
|
|
def __init__(self, service, form_data=None, json_data=None):
|
|
self.app = SimpleNamespace(state=SimpleNamespace(knowledge_upload_service=service))
|
|
self.form_data = form_data
|
|
self.json_data = json_data
|
|
|
|
async def form(self):
|
|
return self.form_data
|
|
|
|
async def json(self):
|
|
return self.json_data
|
|
|
|
|
|
class KnowledgeRouterTests(unittest.IsolatedAsyncioTestCase):
|
|
async def test_preview_endpoint_accepts_file_and_manual_strategy(self):
|
|
with tempfile.TemporaryDirectory() as tmp:
|
|
service = KnowledgeUploadService(
|
|
storage_dir=tmp,
|
|
milvus_client=AsyncMock(),
|
|
embedder=AsyncMock(),
|
|
publisher=AsyncMock(),
|
|
)
|
|
request = FakeRequest(
|
|
service,
|
|
{
|
|
"file": FakeUpload("faq.md", b"Q: Q\nA: A"),
|
|
"strategy": "qa_pair",
|
|
},
|
|
)
|
|
|
|
response = await preview_document_upload(
|
|
request,
|
|
SimpleNamespace(user_type="ADMIN", employee_role=None),
|
|
)
|
|
|
|
self.assertEqual(response.code, 200)
|
|
self.assertEqual(response.data["strategy"], "qa_pair")
|
|
|
|
async def test_confirm_endpoint_returns_ingestion_result(self):
|
|
with tempfile.TemporaryDirectory() as tmp:
|
|
service = KnowledgeUploadService(
|
|
storage_dir=tmp,
|
|
milvus_client=AsyncMock(),
|
|
embedder=AsyncMock(return_value=[[0.0] * 768]),
|
|
publisher=AsyncMock(),
|
|
)
|
|
preview = await service.preview("faq.md", b"Q: Q\nA: A", strategy="qa_pair")
|
|
request = FakeRequest(
|
|
service,
|
|
json_data={
|
|
"upload_id": preview["upload_id"],
|
|
"title": "FAQ",
|
|
"doc_id": "doc-1",
|
|
"collection_name": "fin_faq",
|
|
"strategy": "qa_pair",
|
|
},
|
|
)
|
|
|
|
response = await confirm_document_upload(
|
|
request,
|
|
SimpleNamespace(user_type="ADMIN", employee_role=None),
|
|
)
|
|
|
|
self.assertEqual(response.code, 200)
|
|
self.assertEqual(response.data["doc_id"], "doc-1")
|
|
|
|
async def test_delete_endpoint_returns_deleted_document(self):
|
|
with tempfile.TemporaryDirectory() as tmp:
|
|
service = KnowledgeUploadService(
|
|
storage_dir=tmp,
|
|
milvus_client=AsyncMock(),
|
|
embedder=AsyncMock(),
|
|
publisher=AsyncMock(),
|
|
)
|
|
request = FakeRequest(service, json_data={"doc_id": "doc-delete-1"})
|
|
|
|
response = await delete_document(
|
|
request,
|
|
"doc-delete-1",
|
|
SimpleNamespace(user_type="ADMIN", employee_role=None),
|
|
)
|
|
|
|
self.assertEqual(response.code, 200)
|
|
self.assertEqual(response.data, {"doc_id": "doc-delete-1", "deleted": True})
|
|
|
|
async def test_list_endpoint_returns_documents(self):
|
|
with tempfile.TemporaryDirectory() as tmp:
|
|
service = KnowledgeUploadService(
|
|
storage_dir=tmp,
|
|
milvus_client=AsyncMock(),
|
|
embedder=AsyncMock(),
|
|
publisher=AsyncMock(),
|
|
)
|
|
service.list_documents = AsyncMock(return_value=[{"doc_id": "doc-1"}])
|
|
response = await list_documents(
|
|
FakeRequest(service),
|
|
SimpleNamespace(user_type="ADMIN", employee_role=None),
|
|
)
|
|
|
|
self.assertEqual(response.code, 200)
|
|
self.assertEqual(response.data, [{"doc_id": "doc-1"}])
|
|
|
|
async def test_detail_endpoint_returns_document(self):
|
|
with tempfile.TemporaryDirectory() as tmp:
|
|
service = KnowledgeUploadService(
|
|
storage_dir=tmp,
|
|
milvus_client=AsyncMock(),
|
|
embedder=AsyncMock(),
|
|
publisher=AsyncMock(),
|
|
)
|
|
service.get_document = AsyncMock(return_value={"doc_id": "doc-1"})
|
|
response = await get_document(
|
|
FakeRequest(service),
|
|
"doc-1",
|
|
SimpleNamespace(user_type="ADMIN", employee_role=None),
|
|
)
|
|
|
|
self.assertEqual(response.code, 200)
|
|
self.assertEqual(response.data, {"doc_id": "doc-1"})
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|