43 lines
1.7 KiB
Python
43 lines
1.7 KiB
Python
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()
|