mirror of
https://github.com/OpenHands/OpenHands.git
synced 2025-12-26 05:48:36 +08:00
297 lines
9.5 KiB
Python
297 lines
9.5 KiB
Python
import json
|
|
from unittest.mock import AsyncMock, MagicMock, patch
|
|
|
|
import pytest
|
|
|
|
# Import the module, not the functions directly to avoid circular imports
|
|
import openhands.mcp.utils
|
|
from openhands.core.config.mcp_config import MCPSSEServerConfig, MCPStdioServerConfig
|
|
from openhands.events.action.mcp import MCPAction
|
|
from openhands.events.observation.mcp import MCPObservation
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_create_mcp_clients_empty():
|
|
"""Test creating MCP clients with empty server list."""
|
|
clients = await openhands.mcp.utils.create_mcp_clients([], [])
|
|
assert clients == []
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@patch('openhands.mcp.utils.MCPClient')
|
|
async def test_create_mcp_clients_success(mock_mcp_client):
|
|
"""Test successful creation of MCP clients."""
|
|
# Setup mock
|
|
mock_client_instance = AsyncMock()
|
|
mock_mcp_client.return_value = mock_client_instance
|
|
mock_client_instance.connect_http = AsyncMock()
|
|
|
|
# Test with two servers
|
|
server_configs = [
|
|
MCPSSEServerConfig(url='http://server1:8080'),
|
|
MCPSSEServerConfig(url='http://server2:8080', api_key='test-key'),
|
|
]
|
|
|
|
clients = await openhands.mcp.utils.create_mcp_clients(server_configs, [])
|
|
|
|
# Verify
|
|
assert len(clients) == 2
|
|
assert mock_mcp_client.call_count == 2
|
|
|
|
# Check that connect_http was called with correct parameters
|
|
mock_client_instance.connect_http.assert_any_call(
|
|
server_configs[0], conversation_id=None
|
|
)
|
|
mock_client_instance.connect_http.assert_any_call(
|
|
server_configs[1], conversation_id=None
|
|
)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@patch('openhands.mcp.utils.MCPClient')
|
|
async def test_create_mcp_clients_connection_failure(mock_mcp_client):
|
|
"""Test handling of connection failures when creating MCP clients."""
|
|
# Setup mock
|
|
mock_client_instance = AsyncMock()
|
|
mock_mcp_client.return_value = mock_client_instance
|
|
|
|
# First connection succeeds, second fails
|
|
mock_client_instance.connect_http.side_effect = [
|
|
None, # Success
|
|
Exception('Connection failed'), # Failure
|
|
]
|
|
|
|
server_configs = [
|
|
MCPSSEServerConfig(url='http://server1:8080'),
|
|
MCPSSEServerConfig(url='http://server2:8080'),
|
|
]
|
|
|
|
clients = await openhands.mcp.utils.create_mcp_clients(server_configs, [])
|
|
|
|
# Verify only one client was successfully created
|
|
assert len(clients) == 1
|
|
|
|
|
|
def test_convert_mcp_clients_to_tools_empty():
|
|
"""Test converting empty MCP clients list to tools."""
|
|
tools = openhands.mcp.utils.convert_mcp_clients_to_tools(None)
|
|
assert tools == []
|
|
|
|
tools = openhands.mcp.utils.convert_mcp_clients_to_tools([])
|
|
assert tools == []
|
|
|
|
|
|
def test_convert_mcp_clients_to_tools():
|
|
"""Test converting MCP clients to tools."""
|
|
# Create mock clients with tools
|
|
mock_client1 = MagicMock()
|
|
mock_client2 = MagicMock()
|
|
|
|
# Create mock tools
|
|
mock_tool1 = MagicMock()
|
|
mock_tool1.to_param.return_value = {'function': {'name': 'tool1'}}
|
|
|
|
mock_tool2 = MagicMock()
|
|
mock_tool2.to_param.return_value = {'function': {'name': 'tool2'}}
|
|
|
|
mock_tool3 = MagicMock()
|
|
mock_tool3.to_param.return_value = {'function': {'name': 'tool3'}}
|
|
|
|
# Set up the clients with their tools
|
|
mock_client1.tools = [mock_tool1, mock_tool2]
|
|
mock_client2.tools = [mock_tool3]
|
|
|
|
# Convert to tools
|
|
tools = openhands.mcp.utils.convert_mcp_clients_to_tools(
|
|
[mock_client1, mock_client2]
|
|
)
|
|
|
|
# Verify
|
|
assert len(tools) == 3
|
|
assert tools[0] == {'function': {'name': 'tool1'}}
|
|
assert tools[1] == {'function': {'name': 'tool2'}}
|
|
assert tools[2] == {'function': {'name': 'tool3'}}
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_call_tool_mcp_no_clients():
|
|
"""Test calling MCP tool with no clients."""
|
|
action = MCPAction(name='test_tool', arguments={'arg1': 'value1'})
|
|
|
|
with pytest.raises(ValueError, match='No MCP clients found'):
|
|
await openhands.mcp.utils.call_tool_mcp([], action)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_call_tool_mcp_no_matching_client():
|
|
"""Test calling MCP tool with no matching client."""
|
|
# Create mock client without the requested tool
|
|
mock_client = MagicMock()
|
|
mock_client.tools = [MagicMock(name='other_tool')]
|
|
|
|
action = MCPAction(name='test_tool', arguments={'arg1': 'value1'})
|
|
|
|
with pytest.raises(ValueError, match='No matching MCP agent found for tool name'):
|
|
await openhands.mcp.utils.call_tool_mcp([mock_client], action)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_call_tool_mcp_success():
|
|
"""Test successful MCP tool call."""
|
|
# Create mock client with the requested tool
|
|
mock_client = MagicMock()
|
|
mock_tool = MagicMock()
|
|
# Set the name attribute properly for the tool
|
|
mock_tool.name = 'test_tool'
|
|
mock_client.tools = [mock_tool]
|
|
|
|
# Setup response
|
|
mock_response = MagicMock()
|
|
mock_response.model_dump.return_value = {'result': 'success'}
|
|
|
|
# Setup call_tool method
|
|
mock_client.call_tool = AsyncMock(return_value=mock_response)
|
|
|
|
action = MCPAction(name='test_tool', arguments={'arg1': 'value1'})
|
|
|
|
# Call the function
|
|
observation = await openhands.mcp.utils.call_tool_mcp([mock_client], action)
|
|
|
|
# Verify
|
|
assert isinstance(observation, MCPObservation)
|
|
assert json.loads(observation.content) == {'result': 'success'}
|
|
mock_client.call_tool.assert_called_once_with('test_tool', {'arg1': 'value1'})
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@patch('openhands.mcp.utils.MCPClient')
|
|
async def test_create_mcp_clients_stdio_success(mock_mcp_client):
|
|
"""Test successful creation of MCP clients with stdio servers."""
|
|
# Setup mock
|
|
mock_client_instance = AsyncMock()
|
|
mock_mcp_client.return_value = mock_client_instance
|
|
mock_client_instance.connect_stdio = AsyncMock()
|
|
|
|
# Test with stdio servers
|
|
stdio_server_configs = [
|
|
MCPStdioServerConfig(
|
|
name='test-server-1',
|
|
command='python',
|
|
args=['-m', 'server1'],
|
|
env={'DEBUG': 'true'},
|
|
),
|
|
MCPStdioServerConfig(
|
|
name='test-server-2',
|
|
command='node',
|
|
args=['server2.js'],
|
|
env={'NODE_ENV': 'development'},
|
|
),
|
|
]
|
|
|
|
clients = await openhands.mcp.utils.create_mcp_clients(
|
|
[], [], stdio_servers=stdio_server_configs
|
|
)
|
|
|
|
# Verify
|
|
assert len(clients) == 2
|
|
assert mock_mcp_client.call_count == 2
|
|
|
|
# Check that connect_stdio was called with correct parameters
|
|
mock_client_instance.connect_stdio.assert_any_call(stdio_server_configs[0])
|
|
mock_client_instance.connect_stdio.assert_any_call(stdio_server_configs[1])
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@patch('openhands.mcp.utils.MCPClient')
|
|
async def test_create_mcp_clients_stdio_connection_failure(mock_mcp_client):
|
|
"""Test handling of stdio connection failures when creating MCP clients."""
|
|
# Setup mock
|
|
mock_client_instance = AsyncMock()
|
|
mock_mcp_client.return_value = mock_client_instance
|
|
|
|
# First connection succeeds, second fails
|
|
mock_client_instance.connect_stdio.side_effect = [
|
|
None, # Success
|
|
Exception('Stdio connection failed'), # Failure
|
|
]
|
|
|
|
stdio_server_configs = [
|
|
MCPStdioServerConfig(name='server1', command='python'),
|
|
MCPStdioServerConfig(name='server2', command='invalid_command'),
|
|
]
|
|
|
|
clients = await openhands.mcp.utils.create_mcp_clients(
|
|
[], [], stdio_servers=stdio_server_configs
|
|
)
|
|
|
|
# Verify only one client was successfully created
|
|
assert len(clients) == 1
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@patch('openhands.mcp.utils.create_mcp_clients')
|
|
async def test_fetch_mcp_tools_from_config_with_stdio(mock_create_clients):
|
|
"""Test fetching MCP tools with stdio servers enabled."""
|
|
from openhands.core.config.mcp_config import MCPConfig
|
|
|
|
# Setup mock clients
|
|
mock_client = MagicMock()
|
|
mock_tool = MagicMock()
|
|
mock_tool.to_param.return_value = {'function': {'name': 'stdio_tool'}}
|
|
mock_client.tools = [mock_tool]
|
|
mock_create_clients.return_value = [mock_client]
|
|
|
|
# Create config with stdio servers
|
|
mcp_config = MCPConfig(
|
|
stdio_servers=[MCPStdioServerConfig(name='test-server', command='python')]
|
|
)
|
|
|
|
# Test with use_stdio=True
|
|
tools = await openhands.mcp.utils.fetch_mcp_tools_from_config(
|
|
mcp_config, conversation_id='test-conv', use_stdio=True
|
|
)
|
|
|
|
# Verify
|
|
assert len(tools) == 1
|
|
assert tools[0] == {'function': {'name': 'stdio_tool'}}
|
|
|
|
# Verify create_mcp_clients was called with stdio servers
|
|
mock_create_clients.assert_called_once_with(
|
|
[], [], 'test-conv', mcp_config.stdio_servers
|
|
)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_call_tool_mcp_stdio_client():
|
|
"""Test calling MCP tool on a stdio client."""
|
|
# Create mock stdio client with the requested tool
|
|
mock_client = MagicMock()
|
|
mock_tool = MagicMock()
|
|
mock_tool.name = 'stdio_test_tool'
|
|
mock_client.tools = [mock_tool]
|
|
|
|
# Setup response
|
|
mock_response = MagicMock()
|
|
mock_response.model_dump.return_value = {
|
|
'result': 'stdio_success',
|
|
'data': 'test_data',
|
|
}
|
|
|
|
# Setup call_tool method
|
|
mock_client.call_tool = AsyncMock(return_value=mock_response)
|
|
|
|
action = MCPAction(name='stdio_test_tool', arguments={'input': 'test_input'})
|
|
|
|
# Call the function
|
|
observation = await openhands.mcp.utils.call_tool_mcp([mock_client], action)
|
|
|
|
# Verify
|
|
assert isinstance(observation, MCPObservation)
|
|
assert json.loads(observation.content) == {
|
|
'result': 'stdio_success',
|
|
'data': 'test_data',
|
|
}
|
|
mock_client.call_tool.assert_called_once_with(
|
|
'stdio_test_tool', {'input': 'test_input'}
|
|
)
|