42 lines
1.5 KiB
Python
42 lines
1.5 KiB
Python
import unittest
|
|
from unittest.mock import AsyncMock
|
|
|
|
from rag.generation import generate_answer
|
|
|
|
|
|
class GenerationTests(unittest.IsolatedAsyncioTestCase):
|
|
async def test_switches_to_backup_model_after_primary_failure(self):
|
|
llm = AsyncMock()
|
|
llm.chat.side_effect = [RuntimeError("primary down"), "backup answer"]
|
|
config = {
|
|
"agent.customer.llm.fallback_model": "backup-model",
|
|
"agent.customer.template.system_busy": "系统繁忙，请稍后再试",
|
|
}
|
|
|
|
result = await generate_answer(
|
|
[{"role": "user", "content": "基金是什么"}],
|
|
llm_client=llm, config_getter=config.get, primary_model="primary-model",
|
|
)
|
|
|
|
self.assertEqual(result, "backup answer")
|
|
self.assertEqual(llm.chat.await_args_list[1].kwargs["model"], "backup-model")
|
|
|
|
async def test_returns_configured_system_busy_template_when_models_fail(self):
|
|
llm = AsyncMock()
|
|
llm.chat.side_effect = RuntimeError("down")
|
|
config = {
|
|
"agent.customer.llm.fallback_model": "backup-model",
|
|
"agent.customer.template.system_busy": "系统繁忙，请稍后再试",
|
|
}
|
|
|
|
result = await generate_answer(
|
|
[{"role": "user", "content": "基金是什么"}],
|
|
llm_client=llm, config_getter=config.get, primary_model="primary-model",
|
|
)
|
|
|
|
self.assertEqual(result, "系统繁忙，请稍后再试")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|