AgentCoreService/tests/unit_test.py

#! /usr/bin/env python
# -*- coding: utf-8 -*-
# vim:fenc=utf-8

import pytest
from unittest.mock import Mock, MagicMock

import pqai_agent.abtest.client
from pqai_agent.agent_service import AgentService
from pqai_agent.dialogue_manager import DialogueState, TimeContext
from pqai_agent.message_queue_backend import MemoryQueueBackend
from pqai_agent.mq_message import MessageType, MqMessage, MessageChannel
from pqai_agent.response_type_detector import ResponseTypeDetector
from pqai_agent.user_manager import LocalUserManager
import time
import logging

class TestMessageQueues:
    """测试用消息队列实现"""
    def __init__(self, receive_queue, send_queue, human_queue):
        self.receive_queue = receive_queue
        self.send_queue = send_queue
        self.human_queue = human_queue

@pytest.fixture
def test_env():
    """测试环境初始化"""
    logging.getLogger().setLevel(logging.DEBUG)

    user_manager = LocalUserManager()
    user_relation_manager = Mock()
    user_relation_manager.get_user_tags = Mock(return_value=['AgentTest1'])
    user_relation_manager.list_staff_users = Mock(return_value=[{'staff_id': 'staff_id_0', 'user_id': 'user_id_0'}])

    receive_queue = MemoryQueueBackend()
    send_queue = MemoryQueueBackend()
    human_queue = MemoryQueueBackend()

    queues = TestMessageQueues(receive_queue, send_queue, human_queue)

    # 创建Agent服务实例
    service = AgentService(
        receive_backend=receive_queue,
        send_backend=send_queue,
        human_backend=human_queue,
        user_manager=user_manager,
        user_relation_manager=user_relation_manager
    )
    service.user_profile_extractor.extract_profile_info = Mock(return_value=None)
    service.can_send_to_user = Mock(return_value=True)
    service.start()

    # 替换LLM调用为模拟响应
    service._call_chat_api = Mock(return_value="模拟响应")

    yield service, queues

    service.shutdown(sync=True)
    pqai_agent.abtest.client.get_client().shutdown(blocking=True)

def test_agent_state_change(test_env):
    service, _ = test_env
    agent = service.get_agent_instance('staff_id_0', 'user_id_0')
    assert agent.current_state == DialogueState.INITIALIZED
    assert agent.previous_state == DialogueState.INITIALIZED

    agent.do_state_change(DialogueState.MESSAGE_AGGREGATING)
    assert agent.current_state == DialogueState.MESSAGE_AGGREGATING
    assert agent.previous_state == DialogueState.INITIALIZED

    agent.do_state_change(DialogueState.GREETING)
    assert agent.current_state == DialogueState.GREETING
    assert agent.previous_state == DialogueState.INITIALIZED

    agent.do_state_change(DialogueState.MESSAGE_AGGREGATING)
    assert agent.current_state == DialogueState.MESSAGE_AGGREGATING
    assert agent.previous_state == DialogueState.GREETING
    agent.commit()

    agent.do_state_change(DialogueState.CHITCHAT)
    assert agent.current_state == DialogueState.CHITCHAT
    assert agent.previous_state == DialogueState.GREETING

    agent.rollback_state()
    assert agent.current_state == DialogueState.MESSAGE_AGGREGATING
    assert agent.previous_state == DialogueState.GREETING

    agent.do_state_change(DialogueState.CHITCHAT)
    assert agent.current_state == DialogueState.CHITCHAT
    assert agent.previous_state == DialogueState.GREETING

    agent.do_state_change(DialogueState.MESSAGE_AGGREGATING)
    agent.do_state_change(DialogueState.CHITCHAT)
    assert agent.current_state == DialogueState.CHITCHAT
    assert agent.previous_state == DialogueState.CHITCHAT

    agent.do_state_change(DialogueState.MESSAGE_AGGREGATING)
    agent.commit()

    agent.do_state_change(DialogueState.CHITCHAT)
    agent.rollback_state()
    assert agent.current_state == DialogueState.MESSAGE_AGGREGATING

    agent.rollback_state()
    # no state should be rollback
    assert agent.current_state == DialogueState.MESSAGE_AGGREGATING


def test_response_sanitization(test_env):
    case1 = '[2024-01-01 12:00:00] 你好'
    ret1 = AgentService.sanitize_response(case1)
    assert ret1 == '你好'

    case1 = '2024-01-01 12:00:00 你好'
    ret2 = AgentService.sanitize_response(case1)
    assert ret2 == '你好'

def test_normal_conversation_flow(test_env):
    """测试正常对话流程"""
    service, queues = test_env
    service.get_agent_instance('staff_id_0', "user_id_0").message_aggregation_sec = 0

    # 准备测试消息
    test_msg = MqMessage.build(
        MessageType.TEXT, MessageChannel.CORP_WECHAT,
        'user_id_0', 'staff_id_0', '你好', int(time.time() * 1000))
    queues.receive_queue.produce(test_msg)

    # 处理消息
    message = service.receive_queue.consume()
    if message:
        service.process_single_message(message)

    # 验证响应消息
    sent_msg = queues.send_queue.consume()
    assert sent_msg is not None
    assert sent_msg.receiver == "user_id_0"
    assert "模拟响应" in sent_msg.content

def test_aggregated_conversation_flow(test_env):
    """测试聚合对话流程"""
    service, queues = test_env
    service.get_agent_instance('staff_id_0', "user_id_0").message_aggregation_sec = 1

    # 准备测试消息
    ts_begin = int(time.time() * 1000)
    test_msg = MqMessage.build(
        MessageType.TEXT, MessageChannel.CORP_WECHAT,
        'user_id_0', 'staff_id_0', '你好', ts_begin)
    queues.receive_queue.produce(test_msg)
    test_msg = MqMessage.build(
        MessageType.TEXT, MessageChannel.CORP_WECHAT,
        'user_id_0', 'staff_id_0', '我是老李', ts_begin + 500)
    queues.receive_queue.produce(test_msg)

    # 处理消息
    message = service.receive_queue.consume()
    if message:
        service.process_single_message(message)

    # 验证第一次响应消息
    sent_msg = queues.send_queue.consume()
    assert sent_msg is None

    message = service.receive_queue.consume()
    if message:
        service.process_single_message(message)
    # 验证第二次响应消息
    sent_msg = queues.send_queue.consume()
    assert sent_msg is None

    # 模拟定时器产生空消息触发响应
    service.process_single_message(MqMessage.build(
        MessageType.AGGREGATION_TRIGGER, MessageChannel.CORP_WECHAT,
        'user_id_0', 'staff_id_0', None, ts_begin + 2000
    ))
    # 验证第三次响应消息
    sent_msg = queues.send_queue.consume()
    assert sent_msg is not None
    assert sent_msg.receiver == "user_id_0"
    assert "模拟响应" in sent_msg.content

def test_human_intervention_trigger(test_env):
    """测试触发人工干预"""
    service, queues = test_env
    service.get_agent_instance('staff_id_0', "user_id_0").message_aggregation_sec = 0

    # 准备需要人工干预的消息
    test_msg = MqMessage.build(
        MessageType.TEXT, MessageChannel.CORP_WECHAT,
        "user_id_0", "staff_id_0",
        "我需要帮助！", int(time.time() * 1000)
    )
    queues.receive_queue.produce(test_msg)

    # 处理消息
    message = service.receive_queue.consume()
    if message:
        service.process_single_message(message)

    # 验证人工队列消息
    human_msg = queues.human_queue.consume()
    # 由于相关逻辑未启用，临时关闭该测试
    return
    assert human_msg is not None
    assert human_msg.sender == "user_id_0"
    assert "用户对话需人工介入" in human_msg.content

def test_initiative_conversation(test_env):
    """测试主动发起对话"""
    service, queues = test_env
    service.get_agent_instance('staff_id_0', "user_id_0").message_aggregation_sec = 0
    service._call_chat_api = Mock(return_value="主动发起模拟消息")

    # 设置Agent需要主动发起对话
    agent = service.get_agent_instance('staff_id_0', "user_id_0")
    agent.should_initiate_conversation = Mock(return_value=(True, MagicMock()))
    # 发起对话有时间限制
    agent.get_time_context = Mock(return_value=TimeContext.MORNING)

    service._check_initiative_conversations()

    # 验证主动发起的消息 （由于当前有白名单，无法支持测试）
    sent_msg = queues.send_queue.consume()
    # assert sent_msg is not None
    # assert "主动发起" in sent_msg.content

def test_response_type_detector(test_env):
    case1 = '大哥，那可得提前了解下天气，以便安排行程~我帮您查查明天北京天气？'
    assert ResponseTypeDetector.is_chinese_only(case1) == True
    assert ResponseTypeDetector.if_message_suitable_for_voice(case1) == True
    case2 = 'hi'
    assert ResponseTypeDetector.is_chinese_only(case2) == False
    case3 = '这是链接：http://domain.com'
    assert ResponseTypeDetector.is_chinese_only(case3) == False
    case4 = '大哥，那可得提前了解下天气'
    assert ResponseTypeDetector.if_message_suitable_for_voice(case4) == True