245 lines
9 KiB
Python
245 lines
9 KiB
Python
"""Tests for MCP Client Manager."""
|
|
|
|
import pytest
|
|
from unittest.mock import AsyncMock, patch
|
|
from mcp_forge.mcp.manager import MCPClientManager
|
|
|
|
|
|
@pytest.fixture
|
|
def mock_client1():
|
|
"""Create mock MCP client 1."""
|
|
client = AsyncMock()
|
|
client.name = "client1"
|
|
client.is_connected.return_value = False
|
|
client.connect = AsyncMock()
|
|
client.disconnect = AsyncMock()
|
|
client.list_tools = AsyncMock(return_value=["tool1", "tool2"])
|
|
client.get_tool_schema = AsyncMock(return_value={
|
|
"type": "object",
|
|
"properties": {"param1": {"type": "string"}}
|
|
})
|
|
client.call_tool = AsyncMock(return_value={"result": "success"})
|
|
return client
|
|
|
|
|
|
@pytest.fixture
|
|
def mock_client2():
|
|
"""Create mock MCP client 2."""
|
|
client = AsyncMock()
|
|
client.name = "client2"
|
|
client.is_connected.return_value = False
|
|
client.connect = AsyncMock()
|
|
client.disconnect = AsyncMock()
|
|
client.list_tools = AsyncMock(return_value=["tool3", "tool4"])
|
|
client.get_tool_schema = AsyncMock(return_value={
|
|
"type": "object",
|
|
"properties": {"param2": {"type": "number"}}
|
|
})
|
|
client.call_tool = AsyncMock(return_value={"result": "success2"})
|
|
return client
|
|
|
|
|
|
@pytest.fixture
|
|
def client_config():
|
|
"""Create test client configuration."""
|
|
return {
|
|
"client1": {
|
|
"command": "python",
|
|
"args": ["server1.py"],
|
|
"env": {"KEY1": "value1"}
|
|
},
|
|
"client2": {
|
|
"command": "python",
|
|
"args": ["server2.py"],
|
|
"env": {"KEY2": "value2"}
|
|
}
|
|
}
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_initialize_clients_from_config(mock_client1, mock_client2, client_config):
|
|
"""Test initializing clients from configuration."""
|
|
with patch('mcp_forge.mcp.manager.MCPClientWrapper') as mock_wrapper:
|
|
# Setup mock to return different clients for different configs
|
|
mock_wrapper.side_effect = [mock_client1, mock_client2]
|
|
|
|
manager = MCPClientManager(client_config)
|
|
await manager.initialize()
|
|
|
|
# Verify clients were created with correct config
|
|
assert mock_wrapper.call_count == 2
|
|
|
|
# Verify first client created with correct params
|
|
call1 = mock_wrapper.call_args_list[0]
|
|
assert call1[1]["name"] == "client1"
|
|
assert call1[1]["command"] == "python"
|
|
assert call1[1]["args"] == ["server1.py"]
|
|
assert call1[1]["env"] == {"KEY1": "value1"}
|
|
|
|
# Verify second client created with correct params
|
|
call2 = mock_wrapper.call_args_list[1]
|
|
assert call2[1]["name"] == "client2"
|
|
assert call2[1]["command"] == "python"
|
|
assert call2[1]["args"] == ["server2.py"]
|
|
assert call2[1]["env"] == {"KEY2": "value2"}
|
|
|
|
# Verify clients were connected
|
|
mock_client1.connect.assert_called_once()
|
|
mock_client2.connect.assert_called_once()
|
|
|
|
# Verify tools were listed
|
|
mock_client1.list_tools.assert_called_once()
|
|
mock_client2.list_tools.assert_called_once()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_get_client_for_tool(mock_client1, mock_client2, client_config):
|
|
"""Test getting client that provides a specific tool."""
|
|
with patch('mcp_forge.mcp.manager.MCPClientWrapper') as mock_wrapper:
|
|
mock_wrapper.side_effect = [mock_client1, mock_client2]
|
|
|
|
manager = MCPClientManager(client_config)
|
|
await manager.initialize()
|
|
|
|
# Get client for tool1 (from client1)
|
|
client = await manager.get_client_for_tool("tool1")
|
|
assert client == mock_client1
|
|
|
|
# Get client for tool3 (from client2)
|
|
client = await manager.get_client_for_tool("tool3")
|
|
assert client == mock_client2
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_get_client_for_unknown_tool(mock_client1, mock_client2, client_config):
|
|
"""Test getting client for tool that doesn't exist."""
|
|
with patch('mcp_forge.mcp.manager.MCPClientWrapper') as mock_wrapper:
|
|
mock_wrapper.side_effect = [mock_client1, mock_client2]
|
|
|
|
manager = MCPClientManager(client_config)
|
|
await manager.initialize()
|
|
|
|
# Try to get client for non-existent tool
|
|
with pytest.raises(KeyError, match="Tool 'unknown_tool' not found"):
|
|
await manager.get_client_for_tool("unknown_tool")
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_list_all_tools(mock_client1, mock_client2, client_config):
|
|
"""Test listing all tools across all clients."""
|
|
with patch('mcp_forge.mcp.manager.MCPClientWrapper') as mock_wrapper:
|
|
mock_wrapper.side_effect = [mock_client1, mock_client2]
|
|
|
|
manager = MCPClientManager(client_config)
|
|
await manager.initialize()
|
|
|
|
tools = await manager.list_all_tools()
|
|
|
|
# Should have all tools from both clients
|
|
assert set(tools) == {"tool1", "tool2", "tool3", "tool4"}
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_detect_tool_name_collision():
|
|
"""Test detection of tool name collisions across clients."""
|
|
# Create clients with overlapping tool names
|
|
client1 = AsyncMock()
|
|
client1.name = "client1"
|
|
client1.connect = AsyncMock()
|
|
client1.list_tools = AsyncMock(return_value=["tool1", "tool2"])
|
|
|
|
client2 = AsyncMock()
|
|
client2.name = "client2"
|
|
client2.connect = AsyncMock()
|
|
client2.list_tools = AsyncMock(return_value=["tool2", "tool3"]) # tool2 collision!
|
|
|
|
config = {
|
|
"client1": {"command": "python", "args": ["server1.py"]},
|
|
"client2": {"command": "python", "args": ["server2.py"]}
|
|
}
|
|
|
|
with patch('mcp_forge.mcp.manager.MCPClientWrapper') as mock_wrapper:
|
|
mock_wrapper.side_effect = [client1, client2]
|
|
|
|
manager = MCPClientManager(config)
|
|
|
|
# Should raise ValueError about collision during initialization
|
|
with pytest.raises(ValueError, match="Tool name collision.*tool2.*client1.*client2"):
|
|
await manager.initialize()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_get_tool_schema(mock_client1, mock_client2, client_config):
|
|
"""Test getting tool schema via manager."""
|
|
with patch('mcp_forge.mcp.manager.MCPClientWrapper') as mock_wrapper:
|
|
mock_wrapper.side_effect = [mock_client1, mock_client2]
|
|
|
|
manager = MCPClientManager(client_config)
|
|
await manager.initialize()
|
|
|
|
# Get schema for tool1 (from client1)
|
|
schema = await manager.get_tool_schema("tool1")
|
|
assert schema == {"type": "object", "properties": {"param1": {"type": "string"}}}
|
|
mock_client1.get_tool_schema.assert_called_once_with("tool1")
|
|
|
|
# Get schema for tool3 (from client2)
|
|
schema = await manager.get_tool_schema("tool3")
|
|
assert schema == {"type": "object", "properties": {"param2": {"type": "number"}}}
|
|
mock_client2.get_tool_schema.assert_called_once_with("tool3")
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_call_tool(mock_client1, mock_client2, client_config):
|
|
"""Test calling tool via manager."""
|
|
with patch('mcp_forge.mcp.manager.MCPClientWrapper') as mock_wrapper:
|
|
mock_wrapper.side_effect = [mock_client1, mock_client2]
|
|
|
|
manager = MCPClientManager(client_config)
|
|
await manager.initialize()
|
|
|
|
# Call tool1 (from client1)
|
|
result = await manager.call_tool("tool1", {"param1": "value"})
|
|
assert result == {"result": "success"}
|
|
mock_client1.call_tool.assert_called_once_with("tool1", {"param1": "value"})
|
|
|
|
# Call tool3 (from client2)
|
|
result = await manager.call_tool("tool3", {"param2": 42})
|
|
assert result == {"result": "success2"}
|
|
mock_client2.call_tool.assert_called_once_with("tool3", {"param2": 42})
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_shutdown_all_clients(mock_client1, mock_client2, client_config):
|
|
"""Test shutting down all clients."""
|
|
with patch('mcp_forge.mcp.manager.MCPClientWrapper') as mock_wrapper:
|
|
mock_wrapper.side_effect = [mock_client1, mock_client2]
|
|
|
|
manager = MCPClientManager(client_config)
|
|
await manager.initialize()
|
|
|
|
# Shutdown all clients
|
|
await manager.shutdown()
|
|
|
|
# Verify both clients were disconnected
|
|
mock_client1.disconnect.assert_called_once()
|
|
mock_client2.disconnect.assert_called_once()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_manager_before_initialization():
|
|
"""Test that manager methods fail before initialization."""
|
|
config = {"client1": {"command": "python", "args": ["server.py"]}}
|
|
manager = MCPClientManager(config)
|
|
|
|
# Should raise RuntimeError if not initialized
|
|
with pytest.raises(RuntimeError, match="Manager not initialized"):
|
|
await manager.list_all_tools()
|
|
|
|
with pytest.raises(RuntimeError, match="Manager not initialized"):
|
|
await manager.get_client_for_tool("tool1")
|
|
|
|
with pytest.raises(RuntimeError, match="Manager not initialized"):
|
|
await manager.get_tool_schema("tool1")
|
|
|
|
with pytest.raises(RuntimeError, match="Manager not initialized"):
|
|
await manager.call_tool("tool1", {})
|