mcp-forge/tests/mcp/test_client.py

"""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"]