import asyncio
import unittest
from unittest.mock import patch

from service.customer_agent.bootstrap import build_default_runtime, document_exists_in_milvus


class BootstrapTests(unittest.TestCase):
    def test_document_exists_checks_all_knowledge_collections(self):
        class FakeMilvus:
            def __init__(self):
                self.calls = []

            async def query(self, **kwargs):
                self.calls.append(kwargs)
                return [{"doc_id": "doc-1"}] if kwargs["collection_name"] == "fin_policy" else []

        async def run():
            client = FakeMilvus()
            self.assertTrue(await document_exists_in_milvus(client, "doc-1"))
            self.assertEqual(
                [call["collection_name"] for call in client.calls],
                ["fin_faq", "fin_fund_doc", "fin_policy"],
            )

        asyncio.run(run())

    def test_builds_runtime_from_project_clients_and_mysql_config(self):
        with patch("service.customer_agent.bootstrap.redis_client", return_value="redis"), \
             patch("service.customer_agent.bootstrap.milvus_client", return_value="milvus"), \
             patch("service.customer_agent.bootstrap.llm_client", new="llm"), \
             patch("service.customer_agent.bootstrap.build_anonymous_runtime", return_value="runtime") as builder:
            result = build_default_runtime()

        self.assertEqual(result, "runtime")
        self.assertEqual(builder.call_args.kwargs["redis"], "redis")
        self.assertEqual(builder.call_args.kwargs["milvus_client"], "milvus")
        self.assertEqual(builder.call_args.kwargs["llm_client"], "llm")


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