SmartBizTrip/tests/test_chat_reply.py

import json
import unittest
from unittest.mock import Mock

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

from chat_service import ChatService


class ChatReplyTests(unittest.IsolatedAsyncioTestCase):
    def make_service(self, follow_up):
        service = ChatService.__new__(ChatService)
        payload = {"intents": ["out_of_scope"], "user_queries": {}}
        payload.update(follow_up)
        service.llm = RunnableLambda(
            lambda _: AIMessage(content=json.dumps(payload))
        )
        service.memory = Mock(current_task={})
        service.memory.get_short_term_text.return_value = ""
        service.memory.get_profile_text.return_value = ""
        service.messages = []
        service.conversation_history = ""
        return service

    async def test_empty_out_of_scope_reply_has_fallback(self):
        for payload in ({}, {"follow_up_message": ""},
                        {"follow_up_message": "  "}, {"follow_up_message": None}):
            for streaming in (False, True):
                with self.subTest(payload=payload, streaming=streaming):
                    service = self.make_service(payload)
                    if streaming:
                        reply = "".join([part async for part in service.chat_stream("你好")])
                    else:
                        reply = await service.chat("你好")
                    self.assertIn("企业差旅助手", reply)
                    self.assertEqual(service.messages[-1],
                                     {"role": "assistant", "content": reply})
                    service.memory.add_message.assert_called_with("assistant", reply)

    async def test_existing_reply_is_preserved(self):
        for streaming in (False, True):
            service = self.make_service({"follow_up_message": "你刚才说要去北京。"})
            if streaming:
                reply = "".join([part async for part in service.chat_stream("我要去哪里？")])
            else:
                reply = await service.chat("我要去哪里？")
            self.assertEqual(reply, "你刚才说要去北京。")


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