53 lines
2.0 KiB
Python
53 lines
2.0 KiB
Python
|
import unittest
|
||
|
|
from pathlib import Path
|
||
|
|
from unittest.mock import AsyncMock
|
||
|
|
|
||
|
|
from rag.mock_ingest import MOCK_DOCUMENTS, ingest_mock_documents
|
||
|
|
from rag.milvus_delete import delete_document_vectors
|
||
|
|
|
||
|
|
|
||
|
|
class MockIngestTests(unittest.IsolatedAsyncioTestCase):
|
||
|
|
async def test_mock_documents_are_chunked_embedded_and_inserted(self):
|
||
|
|
milvus = AsyncMock()
|
||
|
|
embedder = AsyncMock()
|
||
|
|
embedder.return_value = [[0.0] * 768]
|
||
|
|
|
||
|
|
await ingest_mock_documents(milvus, embedder=embedder)
|
||
|
|
|
||
|
|
self.assertEqual(len(MOCK_DOCUMENTS), 3)
|
||
|
|
self.assertEqual(embedder.await_count, len(MOCK_DOCUMENTS))
|
||
|
|
self.assertEqual(milvus.insert.await_count, len(MOCK_DOCUMENTS))
|
||
|
|
inserted = [call.kwargs for call in milvus.insert.await_args_list]
|
||
|
|
self.assertEqual(
|
||
|
|
{item["collection_name"] for item in inserted},
|
||
|
|
{item.collection for item in MOCK_DOCUMENTS},
|
||
|
|
)
|
||
|
|
for call_args in inserted:
|
||
|
|
rows = call_args["data"]
|
||
|
|
self.assertTrue(rows)
|
||
|
|
self.assertTrue({"doc_id", "title", "section_title", "text", "strategy", "vector"} <= rows[0].keys())
|
||
|
|
self.assertEqual(len(rows[0]["vector"]), 768)
|
||
|
|
|
||
|
|
def test_mock_markdown_files_are_present(self):
|
||
|
|
for item in MOCK_DOCUMENTS:
|
||
|
|
self.assertTrue(Path(item.path).is_file())
|
||
|
|
self.assertEqual(Path(item.path).suffix, ".md")
|
||
|
|
faq = next(item for item in MOCK_DOCUMENTS if item.collection == "fin_faq")
|
||
|
|
self.assertEqual(
|
||
|
|
sum(line.startswith("Q:") for line in Path(faq.path).read_text(encoding="utf-8").splitlines()),
|
||
|
|
40,
|
||
|
|
)
|
||
|
|
|
||
|
|
async def test_deletes_document_rows_by_doc_id_from_all_knowledge_collections(self):
|
||
|
|
milvus = AsyncMock()
|
||
|
|
|
||
|
|
await delete_document_vectors("doc-123", milvus_client=milvus)
|
||
|
|
|
||
|
|
self.assertEqual(milvus.delete.await_count, 3)
|
||
|
|
for call_args in milvus.delete.await_args_list:
|
||
|
|
self.assertEqual(call_args.kwargs["filter"], 'doc_id == "doc-123"')
|
||
|
|
|
||
|
|
|
||
|
|
if __name__ == "__main__":
|
||
|
|
unittest.main()
|