initial commit after one day coding agent session

tests/mcp/test_client.py

@@ -0,0 +1,249 @@
+"""Tests for MCP Client Wrapper."""
+
+import pytest
+from unittest.mock import Mock, AsyncMock, patch
+from mcp_forge.mcp.client import MCPClientWrapper
+
+
+@pytest.fixture
+def mock_fastmcp_client():
+    """Mock fastmcp Client."""
+    client = AsyncMock()
+    
+    # Create tool mocks with actual string names (not Mock.name)
+    tool1 = Mock()
+    tool1.name = "tool1"
+    tool1.description = "Tool 1"
+    tool1.inputSchema = {"type": "object", "properties": {}}
+    
+    tool2 = Mock()
+    tool2.name = "tool2"
+    tool2.description = "Tool 2"
+    tool2.inputSchema = {"type": "object", "properties": {}}
+    
+    # Mock list_tools to return tool objects
+    client.list_tools = AsyncMock(return_value=Mock(tools=[tool1, tool2]))
+    
+    # Mock call_tool to return result with data
+    client.call_tool = AsyncMock(return_value=Mock(
+        data="result",
+        content=[Mock(text="result")],
+        is_error=False
+    ))
+    
+    # Mock context manager
+    client.__aenter__ = AsyncMock(return_value=client)
+    client.__aexit__ = AsyncMock(return_value=None)
+    
+    return client
+
+
+@pytest.mark.asyncio
+async def test_connect_success(mock_fastmcp_client):
+    """Test successful connection to MCP server."""
+    with patch('mcp_forge.mcp.client.Client', return_value=mock_fastmcp_client):
+        client = MCPClientWrapper(
+            name="test-server",
+            command="python",
+            args=["-m", "test_server"]
+        )
+        
+        assert not client.is_connected()
+        
+        await client.connect()
+        
+        assert client.is_connected()
+        mock_fastmcp_client.__aenter__.assert_called_once()
+
+
+@pytest.mark.asyncio
+async def test_connect_with_env(mock_fastmcp_client):
+    """Test connection with environment variables."""
+    with patch('mcp_forge.mcp.client.Client', return_value=mock_fastmcp_client):
+        env = {"API_KEY": "test123", "DEBUG": "true"}
+        client = MCPClientWrapper(
+            name="test-server",
+            command="python",
+            args=["-m", "test_server"],
+            env=env
+        )
+        
+        await client.connect()
+        
+        assert client.is_connected()
+        assert client.env == env
+
+
+@pytest.mark.asyncio
+async def test_disconnect(mock_fastmcp_client):
+    """Test disconnect from MCP server."""
+    with patch('mcp_forge.mcp.client.Client', return_value=mock_fastmcp_client):
+        client = MCPClientWrapper(
+            name="test-server",
+            command="python",
+            args=["-m", "test_server"]
+        )
+        
+        await client.connect()
+        assert client.is_connected()
+        
+        await client.disconnect()
+        
+        assert not client.is_connected()
+        mock_fastmcp_client.__aexit__.assert_called_once()
+
+
+@pytest.mark.asyncio
+async def test_list_tools(mock_fastmcp_client):
+    """Test listing available tools."""
+    with patch('mcp_forge.mcp.client.Client', return_value=mock_fastmcp_client):
+        client = MCPClientWrapper(
+            name="test-server",
+            command="python",
+            args=["-m", "test_server"]
+        )
+        
+        await client.connect()
+        tools = await client.list_tools()
+        
+        assert tools == ["tool1", "tool2"]
+        mock_fastmcp_client.list_tools.assert_called_once()
+
+
+@pytest.mark.asyncio
+async def test_list_tools_not_connected():
+    """Test list_tools raises error when not connected."""
+    client = MCPClientWrapper(
+        name="test-server",
+        command="python",
+        args=["-m", "test_server"]
+    )
+    
+    with pytest.raises(RuntimeError, match="not connected"):
+        await client.list_tools()
+
+
+@pytest.mark.asyncio
+async def test_get_tool_schema(mock_fastmcp_client):
+    """Test getting tool schema."""
+    with patch('mcp_forge.mcp.client.Client', return_value=mock_fastmcp_client):
+        client = MCPClientWrapper(
+            name="test-server",
+            command="python",
+            args=["-m", "test_server"]
+        )
+        
+        await client.connect()
+        schema = await client.get_tool_schema("tool1")
+        
+        assert schema == {"type": "object", "properties": {}}
+
+
+@pytest.mark.asyncio
+async def test_get_tool_schema_not_found(mock_fastmcp_client):
+    """Test get_tool_schema raises error for unknown tool."""
+    with patch('mcp_forge.mcp.client.Client', return_value=mock_fastmcp_client):
+        client = MCPClientWrapper(
+            name="test-server",
+            command="python",
+            args=["-m", "test_server"]
+        )
+        
+        await client.connect()
+        
+        with pytest.raises(KeyError, match="Tool 'unknown' not found"):
+            await client.get_tool_schema("unknown")
+
+
+@pytest.mark.asyncio
+async def test_call_tool_success(mock_fastmcp_client):
+    """Test successful tool call."""
+    with patch('mcp_forge.mcp.client.Client', return_value=mock_fastmcp_client):
+        client = MCPClientWrapper(
+            name="test-server",
+            command="python",
+            args=["-m", "test_server"]
+        )
+        
+        await client.connect()
+        result = await client.call_tool("tool1", {"param": "value"})
+        
+        assert result == "result"
+        mock_fastmcp_client.call_tool.assert_called_once_with("tool1", {"param": "value"})
+
+
+@pytest.mark.asyncio
+async def test_call_tool_not_connected():
+    """Test call_tool raises error when not connected."""
+    client = MCPClientWrapper(
+        name="test-server",
+        command="python",
+        args=["-m", "test_server"]
+    )
+    
+    with pytest.raises(RuntimeError, match="not connected"):
+        await client.call_tool("tool1", {})
+
+
+@pytest.mark.asyncio
+async def test_call_tool_failure(mock_fastmcp_client):
+    """Test tool call failure handling."""
+    mock_fastmcp_client.call_tool = AsyncMock(side_effect=Exception("Tool failed"))
+    
+    with patch('mcp_forge.mcp.client.Client', return_value=mock_fastmcp_client):
+        client = MCPClientWrapper(
+            name="test-server",
+            command="python",
+            args=["-m", "test_server"]
+        )
+        
+        await client.connect()
+        
+        with pytest.raises(RuntimeError, match="Tool call failed.*Tool failed"):
+            await client.call_tool("tool1", {})
+
+
+@pytest.mark.asyncio
+async def test_connection_failure():
+    """Test connection failure handling."""
+    failing_client = AsyncMock()
+    failing_client.__aenter__ = AsyncMock(side_effect=Exception("Connection failed"))
+    
+    with patch('mcp_forge.mcp.client.Client', return_value=failing_client):
+        client = MCPClientWrapper(
+            name="test-server",
+            command="python",
+            args=["-m", "test_server"]
+        )
+        
+        with pytest.raises(RuntimeError, match="Failed to connect.*Connection failed"):
+            await client.connect()
+        
+        assert not client.is_connected()
+
+
+@pytest.mark.asyncio
+async def test_reconnect(mock_fastmcp_client):
+    """Test reconnection after disconnect."""
+    with patch('mcp_forge.mcp.client.Client', return_value=mock_fastmcp_client):
+        client = MCPClientWrapper(
+            name="test-server",
+            command="python",
+            args=["-m", "test_server"]
+        )
+        
+        # First connection
+        await client.connect()
+        assert client.is_connected()
+        
+        # Disconnect
+        await client.disconnect()
+        assert not client.is_connected()
+        
+        # Reconnect
+        await client.connect()
+        assert client.is_connected()
+        
+        # Should be able to use tools
+        tools = await client.list_tools()
+        assert tools == ["tool1", "tool2"]