mirror of
https://github.com/crewAIInc/crewAI.git
synced 2026-01-09 14:37:59 -05:00
268 lines
9.2 KiB
Python
268 lines
9.2 KiB
Python
"""Mock MCP server implementation for testing."""
|
|
|
|
import asyncio
|
|
import json
|
|
from typing import Any, Dict, List
|
|
from unittest.mock import AsyncMock, Mock
|
|
|
|
|
|
class MockMCPTool:
|
|
"""Mock MCP tool for testing."""
|
|
|
|
def __init__(self, name: str, description: str, input_schema: Dict[str, Any] = None):
|
|
self.name = name
|
|
self.description = description
|
|
self.inputSchema = input_schema or {"type": "object", "properties": {}}
|
|
|
|
|
|
class MockMCPServer:
|
|
"""Mock MCP server for testing various scenarios."""
|
|
|
|
def __init__(self, server_url: str, tools: List[MockMCPTool] = None, behavior: str = "normal"):
|
|
self.server_url = server_url
|
|
self.tools = tools or []
|
|
self.behavior = behavior
|
|
self.call_count = 0
|
|
self.initialize_count = 0
|
|
self.list_tools_count = 0
|
|
|
|
def add_tool(self, name: str, description: str, input_schema: Dict[str, Any] = None):
|
|
"""Add a tool to the mock server."""
|
|
tool = MockMCPTool(name, description, input_schema)
|
|
self.tools.append(tool)
|
|
return tool
|
|
|
|
async def simulate_initialize(self):
|
|
"""Simulate MCP session initialization."""
|
|
self.initialize_count += 1
|
|
|
|
if self.behavior == "slow_init":
|
|
await asyncio.sleep(15) # Exceed connection timeout
|
|
elif self.behavior == "init_error":
|
|
raise Exception("Initialization failed")
|
|
elif self.behavior == "auth_error":
|
|
raise Exception("Authentication failed")
|
|
|
|
async def simulate_list_tools(self):
|
|
"""Simulate MCP tools listing."""
|
|
self.list_tools_count += 1
|
|
|
|
if self.behavior == "slow_list":
|
|
await asyncio.sleep(20) # Exceed discovery timeout
|
|
elif self.behavior == "list_error":
|
|
raise Exception("Failed to list tools")
|
|
elif self.behavior == "json_error":
|
|
raise Exception("JSON parsing error in list_tools")
|
|
|
|
mock_result = Mock()
|
|
mock_result.tools = self.tools
|
|
return mock_result
|
|
|
|
async def simulate_call_tool(self, tool_name: str, arguments: Dict[str, Any]):
|
|
"""Simulate MCP tool execution."""
|
|
self.call_count += 1
|
|
|
|
if self.behavior == "slow_execution":
|
|
await asyncio.sleep(35) # Exceed execution timeout
|
|
elif self.behavior == "execution_error":
|
|
raise Exception("Tool execution failed")
|
|
elif self.behavior == "tool_not_found":
|
|
raise Exception(f"Tool {tool_name} not found")
|
|
|
|
# Find the tool
|
|
tool = next((t for t in self.tools if t.name == tool_name), None)
|
|
if not tool and self.behavior == "normal":
|
|
raise Exception(f"Tool {tool_name} not found")
|
|
|
|
# Create mock successful response
|
|
mock_result = Mock()
|
|
mock_result.content = [Mock(text=f"Result from {tool_name} with args: {arguments}")]
|
|
return mock_result
|
|
|
|
|
|
class MockMCPServerFactory:
|
|
"""Factory for creating various types of mock MCP servers."""
|
|
|
|
@staticmethod
|
|
def create_working_server(server_url: str) -> MockMCPServer:
|
|
"""Create a mock server that works normally."""
|
|
server = MockMCPServer(server_url, behavior="normal")
|
|
server.add_tool("search_tool", "Search for information")
|
|
server.add_tool("analysis_tool", "Analyze data")
|
|
return server
|
|
|
|
@staticmethod
|
|
def create_slow_server(server_url: str, slow_operation: str = "init") -> MockMCPServer:
|
|
"""Create a mock server that is slow for testing timeouts."""
|
|
behavior_map = {
|
|
"init": "slow_init",
|
|
"list": "slow_list",
|
|
"execution": "slow_execution"
|
|
}
|
|
|
|
server = MockMCPServer(server_url, behavior=behavior_map.get(slow_operation, "slow_init"))
|
|
server.add_tool("slow_tool", "A slow tool")
|
|
return server
|
|
|
|
@staticmethod
|
|
def create_failing_server(server_url: str, failure_type: str = "connection") -> MockMCPServer:
|
|
"""Create a mock server that fails in various ways."""
|
|
behavior_map = {
|
|
"connection": "init_error",
|
|
"auth": "auth_error",
|
|
"list": "list_error",
|
|
"json": "json_error",
|
|
"execution": "execution_error",
|
|
"tool_missing": "tool_not_found"
|
|
}
|
|
|
|
server = MockMCPServer(server_url, behavior=behavior_map.get(failure_type, "init_error"))
|
|
if failure_type != "tool_missing":
|
|
server.add_tool("failing_tool", "A tool that fails")
|
|
return server
|
|
|
|
@staticmethod
|
|
def create_exa_like_server(server_url: str) -> MockMCPServer:
|
|
"""Create a mock server that mimics the Exa MCP server."""
|
|
server = MockMCPServer(server_url, behavior="normal")
|
|
server.add_tool(
|
|
"web_search_exa",
|
|
"Search the web using Exa AI - performs real-time web searches and can scrape content from specific URLs",
|
|
{"type": "object", "properties": {"query": {"type": "string"}, "num_results": {"type": "integer"}}}
|
|
)
|
|
server.add_tool(
|
|
"get_code_context_exa",
|
|
"Search and get relevant context for any programming task. Exa-code has the highest quality context",
|
|
{"type": "object", "properties": {"query": {"type": "string"}, "language": {"type": "string"}}}
|
|
)
|
|
return server
|
|
|
|
@staticmethod
|
|
def create_weather_like_server(server_url: str) -> MockMCPServer:
|
|
"""Create a mock server that mimics a weather MCP server."""
|
|
server = MockMCPServer(server_url, behavior="normal")
|
|
server.add_tool(
|
|
"get_current_weather",
|
|
"Get current weather conditions for a location",
|
|
{"type": "object", "properties": {"location": {"type": "string"}}}
|
|
)
|
|
server.add_tool(
|
|
"get_forecast",
|
|
"Get weather forecast for the next 5 days",
|
|
{"type": "object", "properties": {"location": {"type": "string"}, "days": {"type": "integer"}}}
|
|
)
|
|
server.add_tool(
|
|
"get_alerts",
|
|
"Get active weather alerts for a region",
|
|
{"type": "object", "properties": {"region": {"type": "string"}}}
|
|
)
|
|
return server
|
|
|
|
|
|
class MCPServerContextManager:
|
|
"""Context manager for mock MCP servers."""
|
|
|
|
def __init__(self, mock_server: MockMCPServer):
|
|
self.mock_server = mock_server
|
|
|
|
async def __aenter__(self):
|
|
return (None, None, None) # read, write, cleanup
|
|
|
|
async def __aexit__(self, exc_type, exc_val, exc_tb):
|
|
pass
|
|
|
|
|
|
class MCPSessionContextManager:
|
|
"""Context manager for mock MCP sessions."""
|
|
|
|
def __init__(self, mock_server: MockMCPServer):
|
|
self.mock_server = mock_server
|
|
|
|
async def __aenter__(self):
|
|
return MockMCPSession(self.mock_server)
|
|
|
|
async def __aexit__(self, exc_type, exc_val, exc_tb):
|
|
pass
|
|
|
|
|
|
class MockMCPSession:
|
|
"""Mock MCP session for testing."""
|
|
|
|
def __init__(self, mock_server: MockMCPServer):
|
|
self.mock_server = mock_server
|
|
|
|
async def initialize(self):
|
|
"""Mock session initialization."""
|
|
await self.mock_server.simulate_initialize()
|
|
|
|
async def list_tools(self):
|
|
"""Mock tools listing."""
|
|
return await self.mock_server.simulate_list_tools()
|
|
|
|
async def call_tool(self, tool_name: str, arguments: Dict[str, Any]):
|
|
"""Mock tool execution."""
|
|
return await self.mock_server.simulate_call_tool(tool_name, arguments)
|
|
|
|
|
|
def mock_streamablehttp_client(server_url: str, mock_server: MockMCPServer):
|
|
"""Create a mock streamable HTTP client for testing."""
|
|
return MCPServerContextManager(mock_server)
|
|
|
|
|
|
def mock_client_session(read, write, mock_server: MockMCPServer):
|
|
"""Create a mock client session for testing."""
|
|
return MCPSessionContextManager(mock_server)
|
|
|
|
|
|
# Convenience functions for common test scenarios
|
|
|
|
def create_successful_exa_mock():
|
|
"""Create a successful Exa-like mock server."""
|
|
return MockMCPServerFactory.create_exa_like_server("https://mcp.exa.ai/mcp")
|
|
|
|
|
|
def create_failing_connection_mock():
|
|
"""Create a mock server that fails to connect."""
|
|
return MockMCPServerFactory.create_failing_server("https://failing.com/mcp", "connection")
|
|
|
|
|
|
def create_timeout_mock():
|
|
"""Create a mock server that times out."""
|
|
return MockMCPServerFactory.create_slow_server("https://slow.com/mcp", "init")
|
|
|
|
|
|
def create_mixed_servers_scenario():
|
|
"""Create a mixed scenario with working and failing servers."""
|
|
return {
|
|
"working": MockMCPServerFactory.create_working_server("https://working.com/mcp"),
|
|
"failing": MockMCPServerFactory.create_failing_server("https://failing.com/mcp"),
|
|
"slow": MockMCPServerFactory.create_slow_server("https://slow.com/mcp"),
|
|
"auth_fail": MockMCPServerFactory.create_failing_server("https://auth-fail.com/mcp", "auth")
|
|
}
|
|
|
|
|
|
# Pytest fixtures for common mock scenarios
|
|
|
|
@pytest.fixture
|
|
def mock_exa_server():
|
|
"""Provide mock Exa server for tests."""
|
|
return create_successful_exa_mock()
|
|
|
|
|
|
@pytest.fixture
|
|
def mock_failing_server():
|
|
"""Provide mock failing server for tests."""
|
|
return create_failing_connection_mock()
|
|
|
|
|
|
@pytest.fixture
|
|
def mock_slow_server():
|
|
"""Provide mock slow server for tests."""
|
|
return create_timeout_mock()
|
|
|
|
|
|
@pytest.fixture
|
|
def mixed_mock_servers():
|
|
"""Provide mixed mock servers scenario."""
|
|
return create_mixed_servers_scenario()
|