initial commit after one day coding agent session
This commit is contained in:
commit
372af75b90
88 changed files with 22694 additions and 0 deletions
245
tests/mcp/test_manager.py
Normal file
245
tests/mcp/test_manager.py
Normal file
|
|
@ -0,0 +1,245 @@
|
|||
"""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", {})
|
||||
Loading…
Add table
Add a link
Reference in a new issue