mcp-forge/tests/mcp/test_injection.py

"""Tests for Tool Injection Generator."""

import pytest
import ast
from unittest.mock import AsyncMock
from mcp_forge.mcp.injection import ToolInjectionGenerator


@pytest.fixture
def mock_client_manager():
    """Create mock MCP client manager with tools."""
    manager = AsyncMock()
    
    # Tool schemas
    manager.get_tool_schema = AsyncMock(side_effect=lambda tool_name: {
        "read_file": {
            "type": "object",
            "properties": {
                "path": {"type": "string", "description": "File path to read"},
                "encoding": {"type": "string", "description": "File encoding"}
            },
            "required": ["path"]
        },
        "write_file": {
            "type": "object",
            "properties": {
                "path": {"type": "string", "description": "File path to write"},
                "content": {"type": "string", "description": "Content to write"},
                "mode": {"type": "string", "description": "Write mode"}
            },
            "required": ["path", "content"]
        },
        "calculate": {
            "type": "object",
            "properties": {
                "expression": {"type": "string", "description": "Math expression"},
                "precision": {"type": "integer", "description": "Decimal precision"}
            },
            "required": ["expression"]
        }
    }[tool_name])
    
    return manager


@pytest.mark.asyncio
async def test_generate_injection_code_is_valid_python(mock_client_manager):
    """Test that generated code is valid Python."""
    generator = ToolInjectionGenerator(mock_client_manager)
    
    code = await generator.generate_injection_code(
        tool_names=["read_file", "write_file"],
        bridge_socket_path="/tmp/bridge.sock"
    )
    
    # Try to parse the generated code
    try:
        ast.parse(code)
    except SyntaxError as e:
        pytest.fail(f"Generated code has syntax error: {e}")


@pytest.mark.asyncio
async def test_generated_code_includes_bridge_client(mock_client_manager):
    """Test that generated code includes bridge client."""
    generator = ToolInjectionGenerator(mock_client_manager)
    
    code = await generator.generate_injection_code(
        tool_names=["read_file"],
        bridge_socket_path="/tmp/bridge.sock"
    )
    
    # Verify bridge client function is included
    assert "_mcp_call" in code
    assert "socket.socket" in code
    assert "socket.AF_UNIX" in code
    assert "/tmp/bridge.sock" in code


@pytest.mark.asyncio
async def test_generated_code_includes_tool_functions(mock_client_manager):
    """Test that generated code includes wrapper functions for each tool."""
    generator = ToolInjectionGenerator(mock_client_manager)
    
    code = await generator.generate_injection_code(
        tool_names=["read_file", "write_file", "calculate"],
        bridge_socket_path="/tmp/bridge.sock"
    )
    
    # Verify tool functions are defined
    assert "def read_file(" in code
    assert "def write_file(" in code
    assert "def calculate(" in code


@pytest.mark.asyncio
async def test_function_signatures_match_schemas(mock_client_manager):
    """Test that function signatures match tool schemas."""
    generator = ToolInjectionGenerator(mock_client_manager)
    
    code = await generator.generate_injection_code(
        tool_names=["read_file", "write_file"],
        bridge_socket_path="/tmp/bridge.sock"
    )
    
    # read_file has path (required) and encoding (optional)
    assert "def read_file(path: str, encoding: str = None)" in code
    
    # write_file has path, content (required) and mode (optional)
    assert "def write_file(path: str, content: str, mode: str = None)" in code


@pytest.mark.asyncio
async def test_generated_functions_have_docstrings(mock_client_manager):
    """Test that generated functions have docstrings."""
    generator = ToolInjectionGenerator(mock_client_manager)
    
    code = await generator.generate_injection_code(
        tool_names=["read_file"],
        bridge_socket_path="/tmp/bridge.sock"
    )
    
    # Verify docstring is present (contains parameter descriptions)
    assert '"""' in code
    assert "File path to read" in code or "path:" in code.lower()


@pytest.mark.asyncio
async def test_generated_functions_call_bridge(mock_client_manager):
    """Test that generated functions call _mcp_call."""
    generator = ToolInjectionGenerator(mock_client_manager)
    
    code = await generator.generate_injection_code(
        tool_names=["read_file"],
        bridge_socket_path="/tmp/bridge.sock"
    )
    
    # Verify function calls _mcp_call with tool name
    lines = code.split('\n')
    in_read_file = False
    found_call = False
    
    for line in lines:
        if "def read_file(" in line:
            in_read_file = True
        if in_read_file and "_mcp_call" in line and "read_file" in line:
            found_call = True
            break
    
    assert found_call, "Generated function should call _mcp_call with tool name"


@pytest.mark.asyncio
async def test_type_hints_from_schema(mock_client_manager):
    """Test that type hints are generated from schema types."""
    generator = ToolInjectionGenerator(mock_client_manager)
    
    code = await generator.generate_injection_code(
        tool_names=["calculate"],
        bridge_socket_path="/tmp/bridge.sock"
    )
    
    # calculate has string expression and integer precision
    assert "expression: str" in code
    assert "precision: int" in code or "precision: " in code


@pytest.mark.asyncio
async def test_required_vs_optional_parameters(mock_client_manager):
    """Test that required and optional parameters are handled correctly."""
    generator = ToolInjectionGenerator(mock_client_manager)
    
    code = await generator.generate_injection_code(
        tool_names=["read_file"],
        bridge_socket_path="/tmp/bridge.sock"
    )
    
    # path is required (no default), encoding is optional (has default)
    assert "def read_file(path: str, encoding: str = None)" in code or \
           "def read_file(path: str, encoding: " in code


@pytest.mark.asyncio
async def test_generated_code_has_imports(mock_client_manager):
    """Test that generated code includes necessary imports."""
    generator = ToolInjectionGenerator(mock_client_manager)
    
    code = await generator.generate_injection_code(
        tool_names=["read_file"],
        bridge_socket_path="/tmp/bridge.sock"
    )
    
    # Verify imports
    assert "import socket" in code
    assert "import json" in code
    assert "from typing import Any" in code or "typing" in code


@pytest.mark.asyncio
async def test_empty_tool_list(mock_client_manager):
    """Test handling of empty tool list."""
    generator = ToolInjectionGenerator(mock_client_manager)
    
    code = await generator.generate_injection_code(
        tool_names=[],
        bridge_socket_path="/tmp/bridge.sock"
    )
    
    # Should still include bridge client
    assert "_mcp_call" in code
    # But no tool functions
    assert "def read_file(" not in code


@pytest.mark.asyncio
async def test_custom_socket_path(mock_client_manager):
    """Test that custom socket path is used correctly."""
    generator = ToolInjectionGenerator(mock_client_manager)
    
    custom_path = "/custom/path/to/socket.sock"
    code = await generator.generate_injection_code(
        tool_names=["read_file"],
        bridge_socket_path=custom_path
    )
    
    # Verify custom path is in generated code
    assert custom_path in code