import tempfile import unittest import os import time from pathlib import Path from unittest.mock import AsyncMock from rag.embedding import EMBEDDING_DIMENSION from rag.events import KNOWLEDGE_UPDATE_EVENT from service.knowledge_base.upload import KnowledgeUploadService, UploadValidationError class KnowledgeUploadTests(unittest.IsolatedAsyncioTestCase): async def test_preview_saves_upload_and_returns_cleaning_and_chunk_preview(self): with tempfile.TemporaryDirectory() as tmp: service = KnowledgeUploadService( storage_dir=tmp, milvus_client=AsyncMock(), embedder=AsyncMock(), publisher=AsyncMock(), ) result = await service.preview( filename="notice.md", content="\ufeff# Notice\r\n\r\n正文".encode(), strategy="chapter_semantic", ) self.assertTrue(result["upload_id"]) self.assertTrue(result["chunks"]) self.assertTrue(Path(tmp, result["stored_filename"]).is_file()) self.assertIn("cleaning_warnings", result) async def test_confirm_ingests_then_publishes_update_and_removes_temp_file(self): with tempfile.TemporaryDirectory() as tmp: publisher = AsyncMock() milvus = AsyncMock() embedder = AsyncMock(return_value=[[0.0] * EMBEDDING_DIMENSION]) service = KnowledgeUploadService( storage_dir=tmp, milvus_client=milvus, embedder=embedder, publisher=publisher, ) result = await service.confirm( filename="faq.md", content="Q: 什么是基金？\nA: 一种集合投资工具。".encode(), title="FAQ", doc_id="doc-upload-1", collection_name="fin_faq", strategy="qa_pair", ) self.assertEqual(result["doc_id"], "doc-upload-1") publisher.assert_awaited_once() self.assertEqual(publisher.await_args.args[0], KNOWLEDGE_UPDATE_EVENT) self.assertEqual(len(list(Path(tmp).glob("*.md"))), 0) self.assertEqual(len(list(Path(tmp).glob("*.json"))), 0) async def test_confirm_passes_custom_chunk_config_to_ingestion(self): with tempfile.TemporaryDirectory() as tmp: milvus = AsyncMock() embedder = AsyncMock(side_effect=lambda texts: [[0.0] * EMBEDDING_DIMENSION for _ in texts]) service = KnowledgeUploadService( storage_dir=tmp, milvus_client=milvus, embedder=embedder, publisher=AsyncMock(), ) await service.confirm( filename="doc.md", content="一二三四五六七八九十十一十二".encode(), title="Doc", doc_id="doc-config", collection_name="fin_fund_doc", strategy="default", chunk_size=10, chunk_overlap=2, ) self.assertEqual(len(embedder.await_args.args[0]), 2) async def test_confirm_removes_vectors_when_update_event_publish_fails(self): with tempfile.TemporaryDirectory() as tmp: milvus = AsyncMock() publisher = AsyncMock(side_effect=RuntimeError("redis unavailable")) service = KnowledgeUploadService( storage_dir=tmp, milvus_client=milvus, embedder=AsyncMock(return_value=[[0.0] * EMBEDDING_DIMENSION]), publisher=publisher, ) with self.assertRaises(RuntimeError): await service.confirm( filename="policy.md", content=b"policy text", title="Policy", doc_id="doc-event-failure", collection_name="fin_policy", strategy="default", ) self.assertEqual(milvus.delete.await_count, 3) async def test_rejects_unsupported_extension_and_oversized_file(self): with tempfile.TemporaryDirectory() as tmp: service = KnowledgeUploadService( storage_dir=tmp, max_upload_bytes=4, milvus_client=AsyncMock(), embedder=AsyncMock(), publisher=AsyncMock(), ) with self.assertRaises(UploadValidationError): await service.preview("file.exe", b"ok", strategy="default") with self.assertRaises(UploadValidationError): await service.preview("file.md", b"12345", strategy="default") async def test_confirm_rejects_unknown_collection_or_strategy_mismatch(self): with tempfile.TemporaryDirectory() as tmp: service = KnowledgeUploadService( storage_dir=tmp, milvus_client=AsyncMock(), embedder=AsyncMock(), publisher=AsyncMock(), ) with self.assertRaises(UploadValidationError): await service.confirm( filename="faq.md", content=b"Q: Q\nA: A", title="FAQ", doc_id="d1", collection_name="evil_collection", strategy="qa_pair", ) async def test_confirm_rejects_duplicate_doc_id_before_embedding(self): with tempfile.TemporaryDirectory() as tmp: embedder = AsyncMock() service = KnowledgeUploadService( storage_dir=tmp, milvus_client=AsyncMock(), embedder=embedder, publisher=AsyncMock(), document_exists=lambda doc_id: True, ) with self.assertRaises(UploadValidationError): await service.confirm( filename="faq.md", content=b"Q: Q\nA: A", title="FAQ", doc_id="d1", collection_name="fin_faq", strategy="qa_pair", ) embedder.assert_not_awaited() async def test_preview_storage_cleanup_removes_expired_upload_and_manifest(self): with tempfile.TemporaryDirectory() as tmp: service = KnowledgeUploadService( storage_dir=tmp, upload_ttl_seconds=60, milvus_client=AsyncMock(), embedder=AsyncMock(), publisher=AsyncMock(), ) preview = await service.preview("faq.md", b"Q: Q\nA: A", strategy="qa_pair") path = Path(tmp, preview["stored_filename"]) manifest = Path(tmp, f"{preview['upload_id']}.json") old = time.time() - 120 os.utime(path, (old, old)) os.utime(manifest, (old, old)) removed = service.cleanup_expired_uploads(now=time.time()) self.assertEqual(removed, 1) self.assertFalse(path.exists()) self.assertFalse(manifest.exists()) async def test_confirm_rejects_unsupported_extension(self): with tempfile.TemporaryDirectory() as tmp: service = KnowledgeUploadService( storage_dir=tmp, milvus_client=AsyncMock(), embedder=AsyncMock(), publisher=AsyncMock(), ) with self.assertRaises(UploadValidationError): await service.confirm( filename="script.exe", content=b"bad", title="Bad", doc_id="bad-doc", collection_name="fin_faq", strategy="default", ) async def test_delete_document_removes_vectors_and_publishes_update(self): with tempfile.TemporaryDirectory() as tmp: milvus = AsyncMock() publisher = AsyncMock() service = KnowledgeUploadService( storage_dir=tmp, milvus_client=milvus, embedder=AsyncMock(), publisher=publisher, ) result = await service.delete_document("doc-delete-1") self.assertEqual(result, {"doc_id": "doc-delete-1", "deleted": True}) self.assertEqual(milvus.delete.await_count, 3) publisher.assert_awaited_once() self.assertEqual(publisher.await_args.args[0], KNOWLEDGE_UPDATE_EVENT) self.assertEqual(publisher.await_args.args[1]["doc_id"], "doc-delete-1") self.assertEqual(publisher.await_args.args[1]["chunk_count"], 0) self.assertEqual(publisher.await_args.args[1]["action"], "deleted") async def test_list_documents_groups_chunks_across_collections(self): with tempfile.TemporaryDirectory() as tmp: milvus = AsyncMock() async def query(**kwargs): if kwargs["collection_name"] == "fin_faq": return [ { "doc_id": "doc-1", "title": "FAQ", "section_title": "基金基础", "strategy": "qa_pair", }, { "doc_id": "doc-1", "title": "FAQ", "section_title": "购买流程", "strategy": "qa_pair", }, ] return [] milvus.query.side_effect = query service = KnowledgeUploadService( storage_dir=tmp, milvus_client=milvus, embedder=AsyncMock(), publisher=AsyncMock(), ) result = await service.list_documents() self.assertEqual(result, [{ "doc_id": "doc-1", "title": "FAQ", "collection_name": "fin_faq", "strategy": "qa_pair", "chunk_count": 2, }]) async def test_get_document_returns_not_found_when_doc_id_is_missing(self): with tempfile.TemporaryDirectory() as tmp: milvus = AsyncMock() milvus.query.return_value = [] service = KnowledgeUploadService( storage_dir=tmp, milvus_client=milvus, embedder=AsyncMock(), publisher=AsyncMock(), ) self.assertIsNone(await service.get_document("missing-doc")) if __name__ == "__main__": unittest.main()