init
This commit is contained in:
raykkk
255
tests/functionality_mcp_test.py
Normal file
255
tests/functionality_mcp_test.py
Normal file
@@ -0,0 +1,255 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
import os
|
||||
|
||||
"""This module contains utility functions for data processing."""
|
||||
from unittest.mock import AsyncMock, Mock, patch
|
||||
|
||||
import pytest
|
||||
from agentscope.agent import ReActAgent
|
||||
from agentscope.formatter import DashScopeChatFormatter
|
||||
from agentscope.mcp import HttpStatefulClient, HttpStatelessClient
|
||||
from agentscope.message import Msg
|
||||
from agentscope.model import DashScopeChatModel
|
||||
from agentscope.tool import Toolkit
|
||||
from browser_use.functionality.mcp import main
|
||||
from pydantic import BaseModel, Field
|
||||
|
||||
|
||||
class NumberResult(BaseModel):
|
||||
"""A simple number result model for structured output."""
|
||||
|
||||
result: int = Field(description="The result of the calculation")
|
||||
|
||||
|
||||
class TestMCPReActAgent:
|
||||
"""Test suite for MCP ReAct agent functionality"""
|
||||
|
||||
@pytest.fixture
|
||||
def mock_toolkit(self) -> Toolkit:
|
||||
"""Create a mocked Toolkit instance"""
|
||||
return Mock(spec=Toolkit)
|
||||
|
||||
@pytest.fixture
|
||||
def mock_stateful_client(self) -> HttpStatefulClient:
|
||||
"""Create a mocked HttpStatefulClient"""
|
||||
client = Mock(spec=HttpStatefulClient)
|
||||
client.connect = AsyncMock()
|
||||
client.close = AsyncMock()
|
||||
client.get_callable_function = AsyncMock()
|
||||
return client
|
||||
|
||||
@pytest.fixture
|
||||
def mock_stateless_client(self) -> HttpStatelessClient:
|
||||
"""Create a mocked HttpStatelessClient"""
|
||||
client = Mock(spec=HttpStatelessClient)
|
||||
return client
|
||||
|
||||
@pytest.fixture
|
||||
def mock_model(self) -> DashScopeChatModel:
|
||||
"""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) -> DashScopeChatFormatter:
|
||||
"""Create a mocked DashScopeChatFormatter"""
|
||||
return Mock(spec=DashScopeChatFormatter)
|
||||
|
||||
@pytest.fixture
|
||||
def mock_agent(
|
||||
self,
|
||||
mock_model: DashScopeChatModel,
|
||||
mock_formatter: DashScopeChatFormatter,
|
||||
mock_toolkit: Toolkit,
|
||||
) -> Mock:
|
||||
"""Create a mocked ReActAgent instance"""
|
||||
agent = Mock(spec=ReActAgent)
|
||||
agent.model = mock_model
|
||||
agent.formatter = mock_formatter
|
||||
agent.toolkit = mock_toolkit
|
||||
agent.__call__ = AsyncMock(
|
||||
return_value=Mock(
|
||||
metadata={"result": 123456},
|
||||
),
|
||||
)
|
||||
return agent
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_mcp_client_initialization(self) -> None:
|
||||
"""Test MCP client initialization with different transports"""
|
||||
with patch.dict(os.environ, {"DASHSCOPE_API_KEY": "test_key"}):
|
||||
# Test stateful client creation
|
||||
stateful_client = HttpStatefulClient(
|
||||
name="add_client",
|
||||
transport="sse",
|
||||
url="http://localhost:8080",
|
||||
)
|
||||
assert stateful_client.name == "add_client"
|
||||
assert stateful_client.transport == "sse"
|
||||
assert stateful_client.url == "http://localhost:8080"
|
||||
|
||||
# Test stateless client creation
|
||||
stateless_client = HttpStatelessClient(
|
||||
name="multiply_client",
|
||||
transport="streamable_http",
|
||||
url="http://localhost:8081",
|
||||
)
|
||||
assert stateless_client.name == "multiply_client"
|
||||
assert stateless_client.transport == "streamable_http"
|
||||
assert stateless_client.url == "http://localhost:8081"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_toolkit_registration(
|
||||
self,
|
||||
mock_toolkit: Toolkit,
|
||||
mock_stateful_client: HttpStatefulClient,
|
||||
mock_stateless_client: HttpStatelessClient,
|
||||
) -> None:
|
||||
"""Test MCP client registration with toolkit"""
|
||||
with patch.dict(os.environ, {"DASHSCOPE_API_KEY": "test_key"}):
|
||||
# Mock connect and register methods
|
||||
mock_toolkit.register_mcp_client = AsyncMock()
|
||||
|
||||
# Verify registration of both clients
|
||||
await mock_toolkit.register_mcp_client(mock_stateful_client)
|
||||
await mock_toolkit.register_mcp_client(mock_stateless_client)
|
||||
|
||||
assert mock_toolkit.register_mcp_client.call_count == 2
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_agent_initialization(
|
||||
self,
|
||||
mock_model: DashScopeChatModel,
|
||||
mock_formatter: DashScopeChatFormatter,
|
||||
mock_toolkit: Toolkit,
|
||||
) -> None:
|
||||
"""Test ReAct agent initialization"""
|
||||
with patch.dict(os.environ, {"DASHSCOPE_API_KEY": "test_key"}):
|
||||
agent = ReActAgent(
|
||||
name="Jarvis",
|
||||
sys_prompt="You're a helpful assistant named Jarvis.",
|
||||
model=mock_model,
|
||||
formatter=mock_formatter,
|
||||
toolkit=mock_toolkit,
|
||||
)
|
||||
|
||||
assert agent.name == "Jarvis"
|
||||
assert (
|
||||
agent.sys_prompt == "You're a helpful assistant named Jarvis."
|
||||
)
|
||||
assert agent.model == mock_model
|
||||
assert agent.formatter == mock_formatter
|
||||
assert agent.toolkit == mock_toolkit
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_structured_output(
|
||||
self,
|
||||
mock_agent: ReActAgent,
|
||||
) -> None:
|
||||
"""Test structured output handling"""
|
||||
with patch.dict(os.environ, {"DASHSCOPE_API_KEY": "test_key"}):
|
||||
# Create test message
|
||||
test_msg = Msg(
|
||||
"user",
|
||||
"Calculate 2345 multiplied by 3456, then add 4567 to the result,"
|
||||
" what is the final outcome?",
|
||||
"user",
|
||||
)
|
||||
|
||||
# Run agent with structured model
|
||||
result = await mock_agent(test_msg, structured_model=NumberResult)
|
||||
|
||||
# Verify structured output
|
||||
assert isinstance(result, Mock)
|
||||
assert result.metadata["result"] == 123456
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_manual_tool_call(
|
||||
self,
|
||||
mock_stateful_client: HttpStatefulClient,
|
||||
) -> None:
|
||||
"""Test manual tool call functionality"""
|
||||
with patch.dict(os.environ, {"DASHSCOPE_API_KEY": "test_key"}):
|
||||
# Mock callable function
|
||||
mock_callable = AsyncMock(return_value=Mock(content="15"))
|
||||
mock_stateful_client.get_callable_function = AsyncMock(
|
||||
return_value=mock_callable,
|
||||
)
|
||||
|
||||
# Call tool manually
|
||||
tool_function = await mock_stateful_client.get_callable_function(
|
||||
"add",
|
||||
)
|
||||
response = await tool_function(a=5, b=10)
|
||||
|
||||
# Verify tool call
|
||||
mock_stateful_client.get_callable_function.assert_called_once_with(
|
||||
"add",
|
||||
wrap_tool_result=True,
|
||||
)
|
||||
mock_callable.assert_called_once_with(a=5, b=10)
|
||||
assert response.content == "15"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_client_lifecycle(
|
||||
self,
|
||||
mock_stateful_client: HttpStatefulClient,
|
||||
) -> None:
|
||||
"""Test MCP client connection and cleanup"""
|
||||
with patch.dict(os.environ, {"DASHSCOPE_API_KEY": "test_key"}):
|
||||
# Test connection
|
||||
await mock_stateful_client.connect()
|
||||
mock_stateful_client.connect.assert_awaited_once()
|
||||
|
||||
# Test cleanup
|
||||
await mock_stateful_client.close()
|
||||
mock_stateful_client.close.assert_awaited_once()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_full_integration_flow(
|
||||
self,
|
||||
mock_stateful_client: HttpStatefulClient,
|
||||
mock_stateless_client: HttpStatelessClient,
|
||||
mock_toolkit: Toolkit,
|
||||
mock_model: DashScopeChatModel,
|
||||
mock_formatter: DashScopeChatFormatter,
|
||||
) -> None:
|
||||
"""Test full integration flow with mocked dependencies"""
|
||||
with patch.dict(os.environ, {"DASHSCOPE_API_KEY": "test_key"}):
|
||||
# Mock async methods
|
||||
mock_toolkit.register_mcp_client = AsyncMock()
|
||||
mock_stateful_client.connect = AsyncMock()
|
||||
mock_model.call = AsyncMock(
|
||||
return_value=Mock(
|
||||
content="Final answer: 8101807",
|
||||
),
|
||||
)
|
||||
|
||||
# Patch the agent class
|
||||
with patch("main.ReActAgent") as mock_agent_class:
|
||||
mock_agent = Mock()
|
||||
mock_agent.__call__ = AsyncMock(
|
||||
return_value=Mock(
|
||||
metadata={"result": 8101807},
|
||||
),
|
||||
)
|
||||
mock_agent_class.return_value = mock_agent
|
||||
|
||||
# Run the main function
|
||||
await main.main()
|
||||
|
||||
# Verify full flow
|
||||
mock_stateful_client.connect.assert_awaited_once()
|
||||
mock_toolkit.register_mcp_client.assert_any_call(
|
||||
mock_stateful_client,
|
||||
)
|
||||
mock_toolkit.register_mcp_client.assert_any_call(
|
||||
mock_stateless_client,
|
||||
)
|
||||
mock_agent_class.assert_called_once()
|
||||
mock_agent.__call__.assert_called_once()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
pytest.main(["-v", __file__])
|
||||
Reference in New Issue
Block a user