248 lines
8.8 KiB
Python
248 lines
8.8 KiB
Python
# -*- coding: utf-8 -*-
|
|
|
|
# test_manual_plan_example.py
|
|
import os
|
|
import pytest
|
|
import asyncio
|
|
from unittest.mock import AsyncMock, Mock, patch
|
|
from agentscope.agent import ReActAgent, UserAgent
|
|
from agentscope.model import DashScopeChatModel
|
|
from agentscope.tool import Toolkit
|
|
from agentscope.message import Msg
|
|
from agentscope.formatter import DashScopeChatFormatter
|
|
from agentscope.plan import PlanNotebook, SubTask
|
|
from agentscope.tool import (
|
|
execute_shell_command,
|
|
execute_python_code,
|
|
write_text_file,
|
|
insert_text_file,
|
|
view_text_file,
|
|
)
|
|
|
|
# 导入 main.py 中的 main 函数
|
|
from browser_use.functionality.plan.main_manual_plan import main, plan_notebook
|
|
|
|
|
|
class TestManualPlanExample:
|
|
"""Test suite for the manual meta_planner_agent example"""
|
|
|
|
@pytest.fixture
|
|
def mock_toolkit(self):
|
|
"""Create a mocked Toolkit instance"""
|
|
return Mock(spec=Toolkit)
|
|
|
|
@pytest.fixture
|
|
def mock_model(self):
|
|
"""Create a mocked DashScopeChatModel"""
|
|
model = Mock(spec=DashScopeChatModel)
|
|
model.call = AsyncMock(
|
|
return_value=Msg("assistant", "test response", role="assistant"),
|
|
)
|
|
return model
|
|
|
|
@pytest.fixture
|
|
def mock_formatter(self):
|
|
"""Create a mocked DashScopeChatFormatter"""
|
|
return Mock(spec=DashScopeChatFormatter)
|
|
|
|
@pytest.fixture
|
|
def mock_plan_notebook(self):
|
|
"""Create a mocked PlanNotebook instance"""
|
|
return Mock(spec=PlanNotebook)
|
|
|
|
@pytest.fixture
|
|
def mock_agent(
|
|
self,
|
|
mock_model,
|
|
mock_formatter,
|
|
mock_toolkit,
|
|
mock_plan_notebook,
|
|
):
|
|
"""Create a mocked ReActAgent instance"""
|
|
agent = Mock(spec=ReActAgent)
|
|
agent.model = mock_model
|
|
agent.formatter = mock_formatter
|
|
agent.toolkit = mock_toolkit
|
|
agent.plan_notebook = mock_plan_notebook
|
|
agent.__call__ = AsyncMock(
|
|
return_value=Msg("assistant", "test response", role="assistant"),
|
|
)
|
|
return agent
|
|
|
|
@pytest.fixture
|
|
def mock_user(self):
|
|
"""Create a mocked UserAgent instance"""
|
|
user = Mock(spec=UserAgent)
|
|
user.__call__ = AsyncMock(
|
|
return_value=Msg("user", "exit", role="user"),
|
|
)
|
|
return user
|
|
|
|
def test_plan_creation(self):
|
|
"""Test meta_planner_agent creation and subtasks registration"""
|
|
assert plan_notebook.current_plan is not None
|
|
assert (
|
|
plan_notebook.current_plan.name
|
|
== "Comprehensive Report on AgentScope"
|
|
)
|
|
assert len(plan_notebook.current_plan.subtasks) == 4
|
|
|
|
# 验证子任务名称
|
|
subtask_names = [
|
|
subtask.name for subtask in plan_notebook.current_plan.subtasks
|
|
]
|
|
expected_names = [
|
|
"Clone the repository",
|
|
"View the documentation",
|
|
"Study the code",
|
|
"Summarize the findings",
|
|
]
|
|
assert subtask_names == expected_names
|
|
|
|
# 验证子任务描述
|
|
subtask_descriptions = [
|
|
subtask.description
|
|
for subtask in plan_notebook.current_plan.subtasks
|
|
]
|
|
expected_descriptions = [
|
|
"Clone the AgentScope GitHub repository from agentscope-ai/agentscope, and ensure it's the latest version.",
|
|
"View the documentation of AgentScope in the repository.",
|
|
"Study the code of AgentScope, focusing on the core modules and their interactions.",
|
|
"Summarize the findings from the documentation and code study, and write a comprehensive report in markdown format.",
|
|
]
|
|
assert subtask_descriptions == expected_descriptions
|
|
|
|
def test_toolkit_initialization(self):
|
|
"""Test toolkit initialization and tool registration"""
|
|
toolkit = Toolkit()
|
|
# Register all required tools
|
|
toolkit.register_tool_function(execute_shell_command)
|
|
toolkit.register_tool_function(execute_python_code)
|
|
toolkit.register_tool_function(write_text_file)
|
|
toolkit.register_tool_function(insert_text_file)
|
|
toolkit.register_tool_function(view_text_file)
|
|
|
|
# ✅ 通过 hasattr 和 callable 验证工具是否注册成功
|
|
assert hasattr(toolkit, "execute_shell_command")
|
|
assert hasattr(toolkit, "execute_python_code")
|
|
assert hasattr(toolkit, "write_text_file")
|
|
assert hasattr(toolkit, "insert_text_file")
|
|
assert hasattr(toolkit, "view_text_file")
|
|
|
|
assert callable(toolkit.execute_shell_command)
|
|
assert callable(toolkit.execute_python_code)
|
|
assert callable(toolkit.write_text_file)
|
|
assert callable(toolkit.insert_text_file)
|
|
assert callable(toolkit.view_text_file)
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_agent_initialization(
|
|
self,
|
|
mock_model,
|
|
mock_formatter,
|
|
mock_toolkit,
|
|
mock_plan_notebook,
|
|
):
|
|
"""Test ReActAgent initialization"""
|
|
with patch.dict(os.environ, {"DASHSCOPE_API_KEY": "test_key"}):
|
|
agent = ReActAgent(
|
|
name="Friday",
|
|
sys_prompt="You're a helpful assistant named Friday.",
|
|
model=mock_model,
|
|
formatter=mock_formatter,
|
|
toolkit=mock_toolkit,
|
|
plan_notebook=mock_plan_notebook,
|
|
)
|
|
|
|
assert agent.name == "Friday"
|
|
assert (
|
|
agent.sys_prompt == "You're a helpful assistant named Friday."
|
|
)
|
|
assert agent.model == mock_model
|
|
assert agent.formatter == mock_formatter
|
|
assert agent.toolkit == mock_toolkit
|
|
assert agent.plan_notebook == mock_plan_notebook
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_message_loop_exits_on_exit(self, mock_agent, mock_user):
|
|
"""Test the message loop exits when user sends 'exit'"""
|
|
with patch(
|
|
"manual_plan_example.asyncio.sleep",
|
|
) as mock_sleep, patch.dict(
|
|
os.environ,
|
|
{"DASHSCOPE_API_KEY": "test_key"},
|
|
):
|
|
# 避免无限循环
|
|
mock_sleep.side_effect = asyncio.TimeoutError()
|
|
|
|
# 替换 main.py 中的 agent 和 user
|
|
with patch(
|
|
"manual_plan_example.ReActAgent",
|
|
return_value=mock_agent,
|
|
), patch("manual_plan_example.UserAgent", return_value=mock_user):
|
|
try:
|
|
await main()
|
|
except asyncio.TimeoutError:
|
|
pass # 期望的退出方式
|
|
|
|
# ✅ 验证 agent 和 user 被正确调用
|
|
mock_agent.__call__.assert_awaited_once()
|
|
mock_user.__call__.assert_awaited_once()
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_full_message_flow(self, mock_agent, mock_user):
|
|
"""Test the complete message flow between agent and user"""
|
|
with patch.dict(os.environ, {"DASHSCOPE_API_KEY": "test_key"}):
|
|
# 模拟 agent 返回的响应
|
|
mock_agent.__call__ = AsyncMock(
|
|
side_effect=[
|
|
Msg("assistant", "response 1", role="assistant"),
|
|
Msg("assistant", "response 2", role="assistant"),
|
|
],
|
|
)
|
|
|
|
# 模拟 user 返回的响应
|
|
mock_user.__call__ = AsyncMock(
|
|
side_effect=[
|
|
Msg("user", "first message", role="user"),
|
|
Msg("user", "exit", role="user"),
|
|
],
|
|
)
|
|
|
|
# 替换 main.py 中的 agent 和 user
|
|
with patch(
|
|
"manual_plan_example.ReActAgent",
|
|
return_value=mock_agent,
|
|
), patch("manual_plan_example.UserAgent", return_value=mock_user):
|
|
try:
|
|
await main()
|
|
except asyncio.TimeoutError:
|
|
pass # 期望的退出方式
|
|
|
|
# ✅ 验证消息流程
|
|
assert mock_agent.__call__.await_count == 2
|
|
assert mock_user.__call__.await_count == 2
|
|
|
|
# ✅ 验证最终消息是 "exit"
|
|
final_msg = mock_user.__call__.call_args_list[-1][0][0]
|
|
assert final_msg.get_text_content() == "exit"
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_main_runs_without_error(self, mock_agent, mock_user):
|
|
"""Test the main function runs without raising exceptions"""
|
|
with patch.dict(os.environ, {"DASHSCOPE_API_KEY": "test_key"}), patch(
|
|
"manual_plan_example.ReActAgent",
|
|
return_value=mock_agent,
|
|
), patch(
|
|
"manual_plan_example.UserAgent",
|
|
return_value=mock_user,
|
|
), patch(
|
|
"manual_plan_example.asyncio.sleep",
|
|
AsyncMock(),
|
|
):
|
|
# 使用 asyncio.run(main()) 来启动测试
|
|
try:
|
|
await main()
|
|
except Exception as e:
|
|
pytest.fail(f"main() raised an unexpected exception: {e}")
|