init

tests/functionality_agent_plan_test.py

@@ -0,0 +1,206 @@
+# -*- coding: utf-8 -*-
+# test_main.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
+from agentscope.tool import (
+    execute_shell_command,
+    execute_python_code,
+    write_text_file,
+    insert_text_file,
+    view_text_file,
+)
+
+from browser_use.functionality.plan.main_agent_managed_plan import main
+
+
+class TestMainFunctionality:
+    """Test suite for the main.py functionality"""
+
+    @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=Mock(content="test response"))
+        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"""
+        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_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,
+                enable_meta_tool=True,
+                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.enable_meta_tool is True
+            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("main.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("main.ReActAgent", return_value=mock_agent), patch(
+                "main.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("main.ReActAgent", return_value=mock_agent), patch(
+                "main.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(
+            "main.ReActAgent",
+            return_value=mock_agent,
+        ), patch("main.UserAgent", return_value=mock_user), patch(
+            "main.asyncio.sleep",
+            AsyncMock(),
+        ):
+            # 使用 asyncio.run(main()) 来启动测试
+            try:
+                await main()
+            except Exception as e:
+                pytest.fail(f"main() raised an unexpected exception: {e}")