Mutual_Fund/tests/test_milvus_collections.py

import unittest
from unittest.mock import AsyncMock, call
from unittest.mock import patch

from rag.milvus_collections import (
    EMBEDDING_DIMENSION,
    KNOWLEDGE_COLLECTIONS,
    ensure_collections,
)


def _existing_client(dim: int) -> AsyncMock:
    client = AsyncMock()
    client.has_collection.return_value = True
    client.describe_collection.return_value = {
        "fields": [{"name": "vector", "params": {"dim": dim}}]
    }
    return client


class MilvusCollectionTests(unittest.IsolatedAsyncioTestCase):
    async def test_creates_all_project_knowledge_collections(self):
        client = AsyncMock()
        client.has_collection.return_value = False

        await ensure_collections(client)

        self.assertEqual(
            client.has_collection.await_args_list,
            [call(name) for name in KNOWLEDGE_COLLECTIONS],
        )
        self.assertEqual(client.create_collection.await_count, len(KNOWLEDGE_COLLECTIONS))

    async def test_reuses_existing_collections(self):
        client = _existing_client(EMBEDDING_DIMENSION)

        await ensure_collections(client)

        client.create_collection.assert_not_awaited()

    async def test_existing_collection_with_wrong_dimension_raises(self):
        client = _existing_client(EMBEDDING_DIMENSION + 1)

        with self.assertRaises(RuntimeError):
            await ensure_collections(client)

    async def test_collection_schema_contains_metadata_and_1024_vector(self):
        client = AsyncMock()
        client.has_collection.return_value = False

        await ensure_collections(client)

        schema = client.create_collection.await_args.kwargs["schema"]
        fields = {field["name"] for field in schema.to_dict()["fields"]}
        self.assertTrue({"chunk_id", "doc_id", "title", "section_title", "text", "strategy", "vector"} <= fields)
        vector = next(field for field in schema.to_dict()["fields"] if field["name"] == "vector")
        self.assertEqual(vector["params"]["dim"], 1024)

    async def test_default_path_uses_configured_milvus_client(self):
        client = _existing_client(EMBEDDING_DIMENSION)

        with patch("rag.milvus_collections.configured_milvus_client", return_value=client):
            await ensure_collections()

        self.assertEqual(client.has_collection.await_count, len(KNOWLEDGE_COLLECTIONS))


if __name__ == "__main__":
    unittest.main()