New tests for MCP implementation

lib/crewai/tests/mocks/mcp_server_mock.py

@@ -0,0 +1,267 @@
+"""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()