390 lines
11 KiB
Python
390 lines
11 KiB
Python
|
"""Tests for Execute Python Tool."""
|
||
|
|
|
||
|
|
import pytest
|
||
|
|
from unittest.mock import Mock, AsyncMock, patch
|
||
|
|
from mcp.types import Tool, TextContent
|
||
|
|
import json
|
||
|
|
|
||
|
|
from mcp_forge.server.tools.execute_python import ExecutePythonTool
|
||
|
|
from mcp_forge.execution.simple.backend import ExecutionResult
|
||
|
|
|
||
|
|
|
||
|
|
@pytest.fixture
|
||
|
|
def mock_simple_backend():
|
||
|
|
"""Mock simple backend."""
|
||
|
|
backend = Mock()
|
||
|
|
backend.execute = AsyncMock(return_value=ExecutionResult(
|
||
|
|
success=True,
|
||
|
|
stdout="Hello World",
|
||
|
|
stderr="",
|
||
|
|
result="42",
|
||
|
|
execution_time=0.5,
|
||
|
|
exit_code=0
|
||
|
|
))
|
||
|
|
return backend
|
||
|
|
|
||
|
|
|
||
|
|
@pytest.fixture
|
||
|
|
def mock_jupyter_backend():
|
||
|
|
"""Mock jupyter backend."""
|
||
|
|
backend = Mock()
|
||
|
|
backend.execute = AsyncMock(return_value=ExecutionResult(
|
||
|
|
success=True,
|
||
|
|
stdout="Stateful execution",
|
||
|
|
stderr="",
|
||
|
|
result="session-123",
|
||
|
|
execution_time=1.2,
|
||
|
|
exit_code=0
|
||
|
|
))
|
||
|
|
backend.execute_in_session = AsyncMock(return_value=ExecutionResult(
|
||
|
|
success=True,
|
||
|
|
stdout="Using existing session",
|
||
|
|
stderr="",
|
||
|
|
result="[1, 2, 3]",
|
||
|
|
execution_time=0.3,
|
||
|
|
exit_code=0
|
||
|
|
))
|
||
|
|
return backend
|
||
|
|
|
||
|
|
|
||
|
|
@pytest.fixture
|
||
|
|
def mock_client_manager():
|
||
|
|
"""Mock MCP client manager."""
|
||
|
|
manager = Mock()
|
||
|
|
manager.list_all_tools = AsyncMock(return_value=[
|
||
|
|
{"name": "github_search_repos", "description": "Search GitHub repositories"},
|
||
|
|
{"name": "filesystem_read", "description": "Read file contents"}
|
||
|
|
])
|
||
|
|
return manager
|
||
|
|
|
||
|
|
|
||
|
|
@pytest.fixture
|
||
|
|
def mock_bridge_server():
|
||
|
|
"""Mock tool bridge server."""
|
||
|
|
server = Mock()
|
||
|
|
server.socket_path = "/tmp/mcp-bridge.sock"
|
||
|
|
server.is_running = Mock(return_value=True)
|
||
|
|
return server
|
||
|
|
|
||
|
|
|
||
|
|
@pytest.fixture
|
||
|
|
def mock_injection_generator():
|
||
|
|
"""Mock tool injection generator."""
|
||
|
|
generator = Mock()
|
||
|
|
generator.generate_injection_code = Mock(return_value="""
|
||
|
|
# MCP Tool Injection
|
||
|
|
def github_search_repos(**kwargs):
|
||
|
|
import socket
|
||
|
|
# ... tool implementation
|
||
|
|
pass
|
||
|
|
""")
|
||
|
|
return generator
|
||
|
|
|
||
|
|
|
||
|
|
@pytest.fixture
|
||
|
|
def mock_config():
|
||
|
|
"""Mock forge configuration."""
|
||
|
|
config = Mock()
|
||
|
|
config.execution = Mock()
|
||
|
|
config.execution.default_backend = "simple"
|
||
|
|
config.execution.default_timeout = 300
|
||
|
|
config.execution.max_timeout = 1800
|
||
|
|
config.execution.default_memory = "512m"
|
||
|
|
config.execution.max_memory = "2g"
|
||
|
|
return config
|
||
|
|
|
||
|
|
|
||
|
|
@pytest.fixture
|
||
|
|
def execute_python_tool(
|
||
|
|
mock_simple_backend,
|
||
|
|
mock_jupyter_backend,
|
||
|
|
mock_client_manager,
|
||
|
|
mock_bridge_server,
|
||
|
|
mock_injection_generator,
|
||
|
|
mock_config
|
||
|
|
):
|
||
|
|
"""Create ExecutePythonTool instance with mocked dependencies."""
|
||
|
|
return ExecutePythonTool(
|
||
|
|
simple_backend=mock_simple_backend,
|
||
|
|
jupyter_backend=mock_jupyter_backend,
|
||
|
|
client_manager=mock_client_manager,
|
||
|
|
bridge_server=mock_bridge_server,
|
||
|
|
injection_generator=mock_injection_generator,
|
||
|
|
config=mock_config
|
||
|
|
)
|
||
|
|
|
||
|
|
|
||
|
|
@pytest.mark.asyncio
|
||
|
|
async def test_get_tool_definition(execute_python_tool):
|
||
|
|
"""Test that tool definition matches MCP spec."""
|
||
|
|
definition = execute_python_tool.get_tool_definition()
|
||
|
|
|
||
|
|
assert isinstance(definition, Tool)
|
||
|
|
assert definition.name == "execute_python"
|
||
|
|
assert definition.description is not None
|
||
|
|
assert "Execute Python code" in definition.description
|
||
|
|
|
||
|
|
# Verify required schema properties
|
||
|
|
schema = definition.inputSchema
|
||
|
|
assert schema["type"] == "object"
|
||
|
|
assert "code" in schema["properties"]
|
||
|
|
assert "mcp_tools" in schema["properties"]
|
||
|
|
assert "session_id" in schema["properties"]
|
||
|
|
assert "backend" in schema["properties"]
|
||
|
|
assert "timeout" in schema["properties"]
|
||
|
|
assert "custom_image" in schema["properties"]
|
||
|
|
assert "environment" in schema["properties"]
|
||
|
|
assert schema["required"] == ["code"]
|
||
|
|
|
||
|
|
|
||
|
|
@pytest.mark.asyncio
|
||
|
|
async def test_execute_simple_backend_stateless(
|
||
|
|
execute_python_tool,
|
||
|
|
mock_simple_backend
|
||
|
|
):
|
||
|
|
"""Test execution with simple backend (stateless)."""
|
||
|
|
arguments = {
|
||
|
|
"code": "print('Hello World'); 42"
|
||
|
|
}
|
||
|
|
|
||
|
|
result = await execute_python_tool.execute(arguments)
|
||
|
|
|
||
|
|
assert isinstance(result, list)
|
||
|
|
assert len(result) == 1
|
||
|
|
assert isinstance(result[0], TextContent)
|
||
|
|
|
||
|
|
# Parse JSON response
|
||
|
|
response = json.loads(result[0].text)
|
||
|
|
assert response["success"] is True
|
||
|
|
assert response["stdout"] == "Hello World"
|
||
|
|
assert response["result"] == "42"
|
||
|
|
assert response["execution_time"] == 0.5
|
||
|
|
assert "session_id" not in response or response["session_id"] is None
|
||
|
|
|
||
|
|
# Verify simple backend was called
|
||
|
|
mock_simple_backend.execute.assert_called_once()
|
||
|
|
|
||
|
|
|
||
|
|
@pytest.mark.asyncio
|
||
|
|
async def test_execute_jupyter_backend_stateful(
|
||
|
|
execute_python_tool,
|
||
|
|
mock_jupyter_backend
|
||
|
|
):
|
||
|
|
"""Test execution with jupyter backend (stateful)."""
|
||
|
|
arguments = {
|
||
|
|
"code": "x = 42; print('Stateful')",
|
||
|
|
"session_id": "test-session",
|
||
|
|
"backend": "jupyter"
|
||
|
|
}
|
||
|
|
|
||
|
|
result = await execute_python_tool.execute(arguments)
|
||
|
|
|
||
|
|
# Parse JSON response
|
||
|
|
response = json.loads(result[0].text)
|
||
|
|
assert response["success"] is True
|
||
|
|
assert "session_id" in response
|
||
|
|
|
||
|
|
# Verify jupyter backend was called
|
||
|
|
mock_jupyter_backend.execute_in_session.assert_called_once()
|
||
|
|
|
||
|
|
|
||
|
|
@pytest.mark.asyncio
|
||
|
|
async def test_execute_with_mcp_tools(
|
||
|
|
execute_python_tool,
|
||
|
|
mock_simple_backend,
|
||
|
|
mock_injection_generator,
|
||
|
|
mock_bridge_server
|
||
|
|
):
|
||
|
|
"""Test execution with MCP tool injection."""
|
||
|
|
arguments = {
|
||
|
|
"code": "repos = github_search_repos(query='test', max_results=10)",
|
||
|
|
"mcp_tools": ["github_search_repos"]
|
||
|
|
}
|
||
|
|
|
||
|
|
result = await execute_python_tool.execute(arguments)
|
||
|
|
|
||
|
|
# Parse JSON response
|
||
|
|
response = json.loads(result[0].text)
|
||
|
|
assert response["success"] is True
|
||
|
|
assert "available_tools" in response
|
||
|
|
assert "github_search_repos" in response["available_tools"]
|
||
|
|
|
||
|
|
# Verify injection generator was called
|
||
|
|
mock_injection_generator.generate_injection_code.assert_called_once_with(
|
||
|
|
tool_names=["github_search_repos"],
|
||
|
|
socket_path="/tmp/mcp-bridge.sock"
|
||
|
|
)
|
||
|
|
|
||
|
|
|
||
|
|
@pytest.mark.asyncio
|
||
|
|
async def test_execute_with_custom_image(
|
||
|
|
execute_python_tool,
|
||
|
|
mock_simple_backend
|
||
|
|
):
|
||
|
|
"""Test execution with custom image."""
|
||
|
|
arguments = {
|
||
|
|
"code": "import pandas as pd; pd.DataFrame()",
|
||
|
|
"custom_image": "mcp-forge/custom:my-ml-env"
|
||
|
|
}
|
||
|
|
|
||
|
|
_result = await execute_python_tool.execute(arguments)
|
||
|
|
|
||
|
|
# Verify backend was called with custom image
|
||
|
|
call_args = mock_simple_backend.execute.call_args
|
||
|
|
assert call_args is not None
|
||
|
|
# Check that custom_image was passed in execution options
|
||
|
|
assert "image" in call_args.kwargs or "custom_image" in call_args.kwargs
|
||
|
|
|
||
|
|
|
||
|
|
@pytest.mark.asyncio
|
||
|
|
async def test_execute_with_environment_template(
|
||
|
|
execute_python_tool,
|
||
|
|
mock_simple_backend
|
||
|
|
):
|
||
|
|
"""Test execution with environment template."""
|
||
|
|
arguments = {
|
||
|
|
"code": "import numpy as np; np.array([1,2,3])",
|
||
|
|
"environment": "datascience"
|
||
|
|
}
|
||
|
|
|
||
|
|
_result = await execute_python_tool.execute(arguments)
|
||
|
|
|
||
|
|
# Verify backend was called with environment specification
|
||
|
|
mock_simple_backend.execute.assert_called_once()
|
||
|
|
call_args = mock_simple_backend.execute.call_args
|
||
|
|
assert call_args is not None
|
||
|
|
|
||
|
|
|
||
|
|
@pytest.mark.asyncio
|
||
|
|
async def test_execute_with_timeout(
|
||
|
|
execute_python_tool,
|
||
|
|
mock_simple_backend
|
||
|
|
):
|
||
|
|
"""Test execution with custom timeout."""
|
||
|
|
arguments = {
|
||
|
|
"code": "import time; time.sleep(10)",
|
||
|
|
"timeout": 5
|
||
|
|
}
|
||
|
|
|
||
|
|
_result = await execute_python_tool.execute(arguments)
|
||
|
|
|
||
|
|
# Verify timeout was passed to backend
|
||
|
|
call_args = mock_simple_backend.execute.call_args
|
||
|
|
assert call_args is not None
|
||
|
|
|
||
|
|
|
||
|
|
@pytest.mark.asyncio
|
||
|
|
async def test_validate_arguments_missing_code(execute_python_tool):
|
||
|
|
"""Test that validation fails when code is missing."""
|
||
|
|
arguments = {}
|
||
|
|
|
||
|
|
with pytest.raises((ValueError, KeyError)):
|
||
|
|
await execute_python_tool.execute(arguments)
|
||
|
|
|
||
|
|
|
||
|
|
@pytest.mark.asyncio
|
||
|
|
async def test_validate_arguments_invalid_backend(execute_python_tool):
|
||
|
|
"""Test that validation fails for invalid backend."""
|
||
|
|
arguments = {
|
||
|
|
"code": "print('test')",
|
||
|
|
"backend": "invalid"
|
||
|
|
}
|
||
|
|
|
||
|
|
with pytest.raises(ValueError, match="backend"):
|
||
|
|
await execute_python_tool.execute(arguments)
|
||
|
|
|
||
|
|
|
||
|
|
@pytest.mark.asyncio
|
||
|
|
async def test_validate_arguments_invalid_timeout(
|
||
|
|
execute_python_tool,
|
||
|
|
mock_config
|
||
|
|
):
|
||
|
|
"""Test that validation fails for timeout exceeding max."""
|
||
|
|
arguments = {
|
||
|
|
"code": "print('test')",
|
||
|
|
"timeout": 3600 # Exceeds max_timeout of 1800
|
||
|
|
}
|
||
|
|
|
||
|
|
with pytest.raises(ValueError, match="Timeout"):
|
||
|
|
await execute_python_tool.execute(arguments)
|
||
|
|
|
||
|
|
|
||
|
|
@pytest.mark.asyncio
|
||
|
|
async def test_backend_selection_with_session_id(execute_python_tool):
|
||
|
|
"""Test that jupyter backend is selected when session_id provided."""
|
||
|
|
backend = execute_python_tool._select_backend(
|
||
|
|
session_id="test-session",
|
||
|
|
backend=None
|
||
|
|
)
|
||
|
|
|
||
|
|
assert backend == "jupyter"
|
||
|
|
|
||
|
|
|
||
|
|
@pytest.mark.asyncio
|
||
|
|
async def test_backend_selection_default(
|
||
|
|
execute_python_tool,
|
||
|
|
mock_config
|
||
|
|
):
|
||
|
|
"""Test that default backend is used when no session_id."""
|
||
|
|
backend = execute_python_tool._select_backend(
|
||
|
|
session_id=None,
|
||
|
|
backend=None
|
||
|
|
)
|
||
|
|
|
||
|
|
assert backend == "simple"
|
||
|
|
|
||
|
|
|
||
|
|
@pytest.mark.asyncio
|
||
|
|
async def test_backend_selection_explicit(execute_python_tool):
|
||
|
|
"""Test that explicit backend is respected."""
|
||
|
|
backend = execute_python_tool._select_backend(
|
||
|
|
session_id=None,
|
||
|
|
backend="jupyter"
|
||
|
|
)
|
||
|
|
|
||
|
|
assert backend == "jupyter"
|
||
|
|
|
||
|
|
|
||
|
|
@pytest.mark.asyncio
|
||
|
|
async def test_execution_error_handling(
|
||
|
|
execute_python_tool,
|
||
|
|
mock_simple_backend
|
||
|
|
):
|
||
|
|
"""Test that execution errors are handled gracefully."""
|
||
|
|
mock_simple_backend.execute = AsyncMock(return_value=ExecutionResult(
|
||
|
|
success=False,
|
||
|
|
stdout="",
|
||
|
|
stderr="NameError: name 'undefined_variable' is not defined",
|
||
|
|
result=None,
|
||
|
|
execution_time=0.1,
|
||
|
|
exit_code=1
|
||
|
|
))
|
||
|
|
|
||
|
|
arguments = {
|
||
|
|
"code": "print(undefined_variable)"
|
||
|
|
}
|
||
|
|
|
||
|
|
result = await execute_python_tool.execute(arguments)
|
||
|
|
|
||
|
|
# Parse JSON response
|
||
|
|
response = json.loads(result[0].text)
|
||
|
|
assert response["success"] is False
|
||
|
|
assert "NameError" in response["stderr"]
|
||
|
|
|
||
|
|
|
||
|
|
@pytest.mark.asyncio
|
||
|
|
async def test_bridge_server_not_running(
|
||
|
|
execute_python_tool,
|
||
|
|
mock_bridge_server
|
||
|
|
):
|
||
|
|
"""Test that error is raised if bridge server not running."""
|
||
|
|
mock_bridge_server.is_running = Mock(return_value=False)
|
||
|
|
|
||
|
|
arguments = {
|
||
|
|
"code": "repos = github_search_repos(query='test')",
|
||
|
|
"mcp_tools": ["github_search_repos"]
|
||
|
|
}
|
||
|
|
|
||
|
|
with pytest.raises(RuntimeError, match="bridge server"):
|
||
|
|
await execute_python_tool.execute(arguments)
|