SmartBizTrip/tests/test_city_weather_date.py

"""接待建议应将用户指定的日期传给天气助手。"""
import unittest
from types import SimpleNamespace
from unittest.mock import AsyncMock, Mock

from langchain_core.messages import AIMessage
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.runnables import RunnableLambda

from chat_service import ChatService
from memory import ConversationMemory


class CityWeatherDateTests(unittest.IsolatedAsyncioTestCase):
    async def test_preserves_explicit_and_relative_dates(self):
        for question in ["2026年9月26日在北京接待客户", "后天在北京接待客户"]:
            service = ChatService.__new__(ChatService)
            service.memory = ConversationMemory()
            service._city_extract_prompt = ChatPromptTemplate.from_template("{query}")
            service.llm = RunnableLambda(lambda _: AIMessage(content="北京"))
            agent = Mock()
            agent.send_task_async = AsyncMock(return_value=SimpleNamespace(
                status=SimpleNamespace(state="completed"),
                artifacts=[{"parts": [{"text": "北京天气结果"}]}],
            ))
            service.agent_network = Mock()
            service.agent_network.get_agent.return_value = agent

            result = await service._fetch_weather_for_city_service(question)

            self.assertEqual(result, "北京天气结果")
            sent_task = agent.send_task_async.await_args.args[0]
            self.assertIn(question, sent_task.message["content"]["text"])


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