ClaudeFlow ported, needs cleanup
This commit is contained in:
@@ -0,0 +1,398 @@
|
||||
"""Unit tests for CleverClaude MCP (Model Context Protocol) client."""
|
||||
|
||||
from unittest.mock import patch
|
||||
|
||||
import pytest
|
||||
|
||||
from cleverclaude.config.settings import Settings
|
||||
from cleverclaude.mcp.client import MCPClient
|
||||
from cleverclaude.mcp.protocol import MCPError, MCPMessage, MCPProtocol, MCPToolInfo
|
||||
from cleverclaude.mcp.types import MCPToolExecutionResult
|
||||
|
||||
|
||||
@pytest.mark.unit
|
||||
@pytest.mark.async_test
|
||||
class TestMCPClient:
|
||||
"""Test suite for MCPClient class."""
|
||||
|
||||
async def test_initialization(self, test_settings: Settings):
|
||||
"""Test MCPClient initialization."""
|
||||
client = MCPClient(test_settings)
|
||||
|
||||
assert client.config == test_settings
|
||||
assert client.protocol is not None
|
||||
assert client.available_tools == {}
|
||||
assert client.connected_servers == []
|
||||
assert client._initialized is False
|
||||
|
||||
async def test_initialize_client(self, test_settings: Settings):
|
||||
"""Test client initialization process."""
|
||||
client = MCPClient(test_settings)
|
||||
|
||||
with patch.object(client, "_connect_to_servers") as mock_connect:
|
||||
mock_connect.return_value = None
|
||||
|
||||
await client.initialize()
|
||||
|
||||
assert client._initialized is True
|
||||
mock_connect.assert_called_once()
|
||||
|
||||
async def test_connect_to_servers(self, test_settings: Settings):
|
||||
"""Test connection to MCP servers."""
|
||||
client = MCPClient(test_settings)
|
||||
|
||||
mock_server_configs = [
|
||||
{"name": "claude-flow-server", "url": "http://localhost:8001/mcp", "enabled": True},
|
||||
{"name": "neural-server", "url": "http://localhost:8002/mcp", "enabled": True},
|
||||
]
|
||||
|
||||
with patch.object(client, "_connect_server") as mock_connect_server:
|
||||
mock_connect_server.return_value = {
|
||||
"swarm_init": MCPToolInfo(name="swarm_init", description="Initialize swarm"),
|
||||
"agent_spawn": MCPToolInfo(name="agent_spawn", description="Spawn agent"),
|
||||
}
|
||||
|
||||
with patch.object(client.config, "mcp_servers", mock_server_configs):
|
||||
await client._connect_to_servers()
|
||||
|
||||
assert len(client.connected_servers) == 2
|
||||
assert "swarm_init" in client.available_tools
|
||||
assert "agent_spawn" in client.available_tools
|
||||
|
||||
async def test_execute_tool_success(self, test_settings: Settings):
|
||||
"""Test successful tool execution."""
|
||||
client = MCPClient(test_settings)
|
||||
await client.initialize()
|
||||
|
||||
# Mock tool availability
|
||||
client.available_tools["swarm_init"] = MCPToolInfo(
|
||||
name="swarm_init",
|
||||
description="Initialize swarm",
|
||||
parameters={"topology": {"type": "string"}, "maxAgents": {"type": "integer"}},
|
||||
)
|
||||
|
||||
expected_result = {"swarm_id": "swarm_123", "topology": "mesh", "status": "created"}
|
||||
|
||||
with patch.object(client.protocol, "execute_tool") as mock_execute:
|
||||
mock_execute.return_value = MCPToolExecutionResult(success=True, result=expected_result, error=None)
|
||||
|
||||
result = await client.execute_tool("swarm_init", {"topology": "mesh", "maxAgents": 5})
|
||||
|
||||
assert result.success is True
|
||||
assert result.result == expected_result
|
||||
mock_execute.assert_called_once_with("swarm_init", {"topology": "mesh", "maxAgents": 5})
|
||||
|
||||
async def test_execute_tool_not_found(self, test_settings: Settings):
|
||||
"""Test executing non-existent tool."""
|
||||
client = MCPClient(test_settings)
|
||||
await client.initialize()
|
||||
|
||||
with pytest.raises(ValueError, match="Tool 'nonexistent_tool' not found"):
|
||||
await client.execute_tool("nonexistent_tool", {})
|
||||
|
||||
async def test_execute_tool_with_error(self, test_settings: Settings):
|
||||
"""Test tool execution with error."""
|
||||
client = MCPClient(test_settings)
|
||||
await client.initialize()
|
||||
|
||||
# Mock tool availability
|
||||
client.available_tools["failing_tool"] = MCPToolInfo(name="failing_tool", description="Tool that fails")
|
||||
|
||||
with patch.object(client.protocol, "execute_tool") as mock_execute:
|
||||
mock_execute.return_value = MCPToolExecutionResult(
|
||||
success=False, result=None, error="Tool execution failed"
|
||||
)
|
||||
|
||||
result = await client.execute_tool("failing_tool", {})
|
||||
|
||||
assert result.success is False
|
||||
assert result.error == "Tool execution failed"
|
||||
|
||||
async def test_get_available_tools(self, test_settings: Settings):
|
||||
"""Test getting list of available tools."""
|
||||
client = MCPClient(test_settings)
|
||||
await client.initialize()
|
||||
|
||||
# Add mock tools
|
||||
client.available_tools = {
|
||||
"swarm_init": MCPToolInfo(name="swarm_init", description="Initialize swarm"),
|
||||
"agent_spawn": MCPToolInfo(name="agent_spawn", description="Spawn agent"),
|
||||
"task_orchestrate": MCPToolInfo(name="task_orchestrate", description="Orchestrate task"),
|
||||
}
|
||||
|
||||
tools = await client.get_available_tools()
|
||||
|
||||
assert len(tools) == 3
|
||||
assert "swarm_init" in tools
|
||||
assert "agent_spawn" in tools
|
||||
assert "task_orchestrate" in tools
|
||||
|
||||
async def test_get_tool_info(self, test_settings: Settings):
|
||||
"""Test getting tool information."""
|
||||
client = MCPClient(test_settings)
|
||||
await client.initialize()
|
||||
|
||||
tool_info = MCPToolInfo(
|
||||
name="swarm_init",
|
||||
description="Initialize swarm topology",
|
||||
parameters={
|
||||
"topology": {"type": "string", "enum": ["mesh", "hierarchical", "star", "ring"]},
|
||||
"maxAgents": {"type": "integer", "minimum": 1, "maximum": 100},
|
||||
},
|
||||
)
|
||||
client.available_tools["swarm_init"] = tool_info
|
||||
|
||||
retrieved_info = await client.get_tool_info("swarm_init")
|
||||
|
||||
assert retrieved_info == tool_info
|
||||
assert retrieved_info.name == "swarm_init"
|
||||
assert retrieved_info.description == "Initialize swarm topology"
|
||||
assert "topology" in retrieved_info.parameters
|
||||
assert "maxAgents" in retrieved_info.parameters
|
||||
|
||||
async def test_execute_multiple_tools(self, test_settings: Settings):
|
||||
"""Test executing multiple tools in sequence."""
|
||||
client = MCPClient(test_settings)
|
||||
await client.initialize()
|
||||
|
||||
# Mock multiple tools
|
||||
tools = ["swarm_init", "agent_spawn", "task_orchestrate"]
|
||||
for tool in tools:
|
||||
client.available_tools[tool] = MCPToolInfo(name=tool, description=f"Execute {tool}")
|
||||
|
||||
execution_sequence = [
|
||||
("swarm_init", {"topology": "mesh"}),
|
||||
("agent_spawn", {"type": "researcher"}),
|
||||
("task_orchestrate", {"task": "test task"}),
|
||||
]
|
||||
|
||||
results = []
|
||||
with patch.object(client.protocol, "execute_tool") as mock_execute:
|
||||
mock_execute.side_effect = [
|
||||
MCPToolExecutionResult(success=True, result={"swarm_id": "swarm_1"}),
|
||||
MCPToolExecutionResult(success=True, result={"agent_id": "agent_1"}),
|
||||
MCPToolExecutionResult(success=True, result={"task_id": "task_1"}),
|
||||
]
|
||||
|
||||
for tool_name, params in execution_sequence:
|
||||
result = await client.execute_tool(tool_name, params)
|
||||
results.append(result)
|
||||
|
||||
assert len(results) == 3
|
||||
assert all(r.success for r in results)
|
||||
assert results[0].result["swarm_id"] == "swarm_1"
|
||||
assert results[1].result["agent_id"] == "agent_1"
|
||||
assert results[2].result["task_id"] == "task_1"
|
||||
|
||||
async def test_batch_execute_tools(self, test_settings: Settings):
|
||||
"""Test batch execution of tools."""
|
||||
client = MCPClient(test_settings)
|
||||
await client.initialize()
|
||||
|
||||
# Mock tools
|
||||
for tool in ["tool_1", "tool_2", "tool_3"]:
|
||||
client.available_tools[tool] = MCPToolInfo(name=tool, description=f"Tool {tool}")
|
||||
|
||||
batch_requests = [
|
||||
{"tool": "tool_1", "params": {"param": "value1"}},
|
||||
{"tool": "tool_2", "params": {"param": "value2"}},
|
||||
{"tool": "tool_3", "params": {"param": "value3"}},
|
||||
]
|
||||
|
||||
with patch.object(client.protocol, "execute_tool") as mock_execute:
|
||||
mock_execute.side_effect = [
|
||||
MCPToolExecutionResult(success=True, result={"result": f"result_{i}"}) for i in range(3)
|
||||
]
|
||||
|
||||
results = await client.batch_execute(batch_requests)
|
||||
|
||||
assert len(results) == 3
|
||||
assert all(r.success for r in results)
|
||||
assert mock_execute.call_count == 3
|
||||
|
||||
async def test_health_check(self, test_settings: Settings):
|
||||
"""Test MCP client health check."""
|
||||
client = MCPClient(test_settings)
|
||||
await client.initialize()
|
||||
|
||||
client.connected_servers = ["server_1", "server_2"]
|
||||
client.available_tools = {
|
||||
"tool_1": MCPToolInfo(name="tool_1", description="Tool 1"),
|
||||
"tool_2": MCPToolInfo(name="tool_2", description="Tool 2"),
|
||||
}
|
||||
|
||||
with patch.object(client.protocol, "ping_server") as mock_ping:
|
||||
mock_ping.return_value = True
|
||||
|
||||
health = await client.health_check()
|
||||
|
||||
assert health["status"] == "healthy"
|
||||
assert health["connected_servers"] == 2
|
||||
assert health["available_tools"] == 2
|
||||
assert health["server_connectivity"]["server_1"] is True
|
||||
assert health["server_connectivity"]["server_2"] is True
|
||||
|
||||
async def test_disconnect(self, test_settings: Settings):
|
||||
"""Test client disconnection and cleanup."""
|
||||
client = MCPClient(test_settings)
|
||||
await client.initialize()
|
||||
|
||||
client.connected_servers = ["server_1", "server_2"]
|
||||
|
||||
with patch.object(client.protocol, "disconnect") as mock_disconnect:
|
||||
await client.disconnect()
|
||||
|
||||
assert len(client.connected_servers) == 0
|
||||
assert len(client.available_tools) == 0
|
||||
assert client._initialized is False
|
||||
mock_disconnect.assert_called_once()
|
||||
|
||||
|
||||
@pytest.mark.unit
|
||||
@pytest.mark.async_test
|
||||
class TestMCPProtocol:
|
||||
"""Test suite for MCPProtocol class."""
|
||||
|
||||
def test_create_message(self):
|
||||
"""Test MCP message creation."""
|
||||
protocol = MCPProtocol()
|
||||
|
||||
message = protocol.create_message(
|
||||
method="tools/execute", params={"tool": "swarm_init", "arguments": {"topology": "mesh"}}
|
||||
)
|
||||
|
||||
assert isinstance(message, MCPMessage)
|
||||
assert message.method == "tools/execute"
|
||||
assert message.params["tool"] == "swarm_init"
|
||||
assert "id" in message.dict() # Should have generated ID
|
||||
|
||||
def test_parse_response_success(self):
|
||||
"""Test parsing successful MCP response."""
|
||||
protocol = MCPProtocol()
|
||||
|
||||
response_data = {
|
||||
"id": "request_123",
|
||||
"result": {"success": True, "data": {"swarm_id": "swarm_456", "status": "created"}},
|
||||
}
|
||||
|
||||
result = protocol.parse_response(response_data)
|
||||
|
||||
assert result["success"] is True
|
||||
assert result["data"]["swarm_id"] == "swarm_456"
|
||||
|
||||
def test_parse_response_error(self):
|
||||
"""Test parsing MCP error response."""
|
||||
protocol = MCPProtocol()
|
||||
|
||||
response_data = {
|
||||
"id": "request_123",
|
||||
"error": {"code": -32601, "message": "Method not found", "data": {"method": "unknown_method"}},
|
||||
}
|
||||
|
||||
with pytest.raises(MCPError, match="Method not found"):
|
||||
protocol.parse_response(response_data)
|
||||
|
||||
def test_validate_tool_parameters(self):
|
||||
"""Test tool parameter validation."""
|
||||
protocol = MCPProtocol()
|
||||
|
||||
tool_info = MCPToolInfo(
|
||||
name="swarm_init",
|
||||
description="Initialize swarm",
|
||||
parameters={
|
||||
"topology": {"type": "string", "enum": ["mesh", "hierarchical"]},
|
||||
"maxAgents": {"type": "integer", "minimum": 1, "maximum": 100},
|
||||
},
|
||||
)
|
||||
|
||||
# Valid parameters
|
||||
valid_params = {"topology": "mesh", "maxAgents": 10}
|
||||
assert protocol.validate_parameters(tool_info, valid_params) is True
|
||||
|
||||
# Invalid topology
|
||||
invalid_params = {"topology": "invalid", "maxAgents": 10}
|
||||
assert protocol.validate_parameters(tool_info, invalid_params) is False
|
||||
|
||||
# Out of range maxAgents
|
||||
invalid_params = {"topology": "mesh", "maxAgents": 200}
|
||||
assert protocol.validate_parameters(tool_info, invalid_params) is False
|
||||
|
||||
async def test_execute_tool_with_retry(self):
|
||||
"""Test tool execution with retry logic."""
|
||||
protocol = MCPProtocol()
|
||||
|
||||
# Mock connection that fails first time, succeeds second time
|
||||
call_count = 0
|
||||
|
||||
async def mock_send_request(message):
|
||||
nonlocal call_count
|
||||
call_count += 1
|
||||
if call_count == 1:
|
||||
raise ConnectionError("Network error")
|
||||
return {"id": message.id, "result": {"success": True, "data": "success"}}
|
||||
|
||||
with patch.object(protocol, "_send_request", side_effect=mock_send_request):
|
||||
result = await protocol.execute_tool("test_tool", {}, max_retries=2)
|
||||
|
||||
assert result.success is True
|
||||
assert call_count == 2 # Should have retried once
|
||||
|
||||
|
||||
@pytest.mark.unit
|
||||
class TestMCPTypes:
|
||||
"""Test suite for MCP type definitions."""
|
||||
|
||||
def test_mcp_message_creation(self):
|
||||
"""Test MCPMessage creation."""
|
||||
message = MCPMessage(method="tools/list", params={"category": "swarm"}, id="msg_123")
|
||||
|
||||
assert message.method == "tools/list"
|
||||
assert message.params == {"category": "swarm"}
|
||||
assert message.id == "msg_123"
|
||||
|
||||
def test_mcp_tool_info_creation(self):
|
||||
"""Test MCPToolInfo creation."""
|
||||
tool_info = MCPToolInfo(
|
||||
name="agent_spawn",
|
||||
description="Spawn a new agent",
|
||||
parameters={
|
||||
"type": {"type": "string", "enum": ["researcher", "coder", "analyst"]},
|
||||
"name": {"type": "string"},
|
||||
"capabilities": {"type": "array", "items": {"type": "string"}},
|
||||
},
|
||||
)
|
||||
|
||||
assert tool_info.name == "agent_spawn"
|
||||
assert tool_info.description == "Spawn a new agent"
|
||||
assert "type" in tool_info.parameters
|
||||
assert "capabilities" in tool_info.parameters
|
||||
|
||||
def test_mcp_tool_execution_result(self):
|
||||
"""Test MCPToolExecutionResult creation."""
|
||||
# Successful result
|
||||
success_result = MCPToolExecutionResult(
|
||||
success=True, result={"agent_id": "agent_123", "status": "active"}, error=None, execution_time=0.5
|
||||
)
|
||||
|
||||
assert success_result.success is True
|
||||
assert success_result.result["agent_id"] == "agent_123"
|
||||
assert success_result.error is None
|
||||
assert success_result.execution_time == 0.5
|
||||
|
||||
# Error result
|
||||
error_result = MCPToolExecutionResult(
|
||||
success=False, result=None, error="Agent creation failed", execution_time=0.1
|
||||
)
|
||||
|
||||
assert error_result.success is False
|
||||
assert error_result.result is None
|
||||
assert error_result.error == "Agent creation failed"
|
||||
|
||||
def test_mcp_error_creation(self):
|
||||
"""Test MCPError exception creation."""
|
||||
error = MCPError(code=-32601, message="Method not found", data={"method": "unknown_method"})
|
||||
|
||||
assert error.code == -32601
|
||||
assert error.message == "Method not found"
|
||||
assert error.data["method"] == "unknown_method"
|
||||
assert "Method not found" in str(error)
|
||||
Reference in New Issue
Block a user