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()