299 lines
10 KiB
Python
299 lines
10 KiB
Python
"""Tests for the Code Executor module."""
|
|
|
|
import pytest
|
|
from unittest.mock import Mock
|
|
import json
|
|
|
|
from mcp_forge.execution.simple.executor import CodeExecutor, ExecutionResult
|
|
from mcp_forge.podman.containers import SecureContainerManager
|
|
from mcp_forge.security.resource_limits import ResourceLimits
|
|
|
|
|
|
@pytest.fixture
|
|
def resource_limits():
|
|
"""Standard resource limits for testing."""
|
|
return ResourceLimits(
|
|
memory="512m",
|
|
cpu_quota=50000,
|
|
storage="1g",
|
|
timeout=30
|
|
)
|
|
|
|
|
|
@pytest.fixture
|
|
def mock_container_manager():
|
|
"""Mock SecureContainerManager."""
|
|
manager = Mock(spec=SecureContainerManager)
|
|
|
|
# Mock container lifecycle
|
|
manager.create_container.return_value = "test-container-123"
|
|
manager.start_container.return_value = None
|
|
manager.stop_container.return_value = None
|
|
manager.remove_container.return_value = None
|
|
manager.wait_for_container.return_value = 0 # exit code
|
|
manager.get_container_logs.return_value = ("", "") # (stdout, stderr)
|
|
|
|
return manager
|
|
|
|
|
|
@pytest.fixture
|
|
def executor(mock_container_manager, resource_limits):
|
|
"""CodeExecutor instance with mocked dependencies."""
|
|
return CodeExecutor(
|
|
container_manager=mock_container_manager,
|
|
image="mcp-forge/python:3.11",
|
|
resource_limits=resource_limits
|
|
)
|
|
|
|
|
|
def test_execute_simple_python_code_returns_result(executor, mock_container_manager):
|
|
"""Test executing simple Python code returns the result."""
|
|
# Mock successful execution with result
|
|
mock_container_manager.get_container_logs.return_value = (
|
|
json.dumps({"result": 42, "error": None}),
|
|
""
|
|
)
|
|
|
|
result = executor.execute("2 + 2")
|
|
|
|
assert result.success is True
|
|
assert result.result == 42
|
|
assert result.exit_code == 0
|
|
assert result.error is None
|
|
|
|
# Verify container lifecycle
|
|
mock_container_manager.create_container.assert_called_once()
|
|
mock_container_manager.start_container.assert_called_once_with("test-container-123")
|
|
mock_container_manager.wait_for_container.assert_called_once_with("test-container-123", timeout=30)
|
|
mock_container_manager.remove_container.assert_called_once_with("test-container-123")
|
|
|
|
|
|
def test_execute_code_with_stdout_capture(executor, mock_container_manager):
|
|
"""Test code execution captures stdout."""
|
|
mock_container_manager.get_container_logs.return_value = (
|
|
json.dumps({"result": None, "error": None}) + "\n" + "Hello, World!",
|
|
""
|
|
)
|
|
|
|
result = executor.execute('print("Hello, World!")')
|
|
|
|
assert result.success is True
|
|
assert "Hello, World!" in result.stdout
|
|
assert result.stderr == ""
|
|
|
|
|
|
def test_execute_code_with_stderr_capture(executor, mock_container_manager):
|
|
"""Test code execution captures stderr."""
|
|
mock_container_manager.get_container_logs.return_value = (
|
|
json.dumps({"result": None, "error": None}),
|
|
"Warning: something happened"
|
|
)
|
|
|
|
result = executor.execute('import sys; print("warning", file=sys.stderr)')
|
|
|
|
assert result.success is True
|
|
assert result.stderr == "Warning: something happened"
|
|
|
|
|
|
def test_execute_code_timeout_enforcement(executor, mock_container_manager):
|
|
"""Test code execution enforces timeout."""
|
|
# Simulate timeout by having wait_for_container take too long
|
|
mock_container_manager.wait_for_container.side_effect = TimeoutError("Container exceeded timeout")
|
|
|
|
result = executor.execute("import time; time.sleep(60)", timeout=1)
|
|
|
|
assert result.success is False
|
|
assert result.error is not None
|
|
assert "timeout" in result.error.lower()
|
|
|
|
# Verify cleanup still happens
|
|
mock_container_manager.remove_container.assert_called_once_with("test-container-123")
|
|
|
|
|
|
def test_execute_code_with_exception_handling(executor, mock_container_manager):
|
|
"""Test code execution handles exceptions gracefully."""
|
|
error_msg = "ZeroDivisionError: division by zero"
|
|
mock_container_manager.get_container_logs.return_value = (
|
|
json.dumps({"result": None, "error": error_msg}),
|
|
""
|
|
)
|
|
mock_container_manager.wait_for_container.return_value = 1 # non-zero exit
|
|
|
|
result = executor.execute("1 / 0")
|
|
|
|
assert result.success is False
|
|
assert result.error == error_msg
|
|
assert result.exit_code == 1
|
|
|
|
|
|
def test_execute_code_with_syntax_error_returns_clear_error(executor, mock_container_manager):
|
|
"""Test code with syntax error returns clear error message."""
|
|
error_msg = "SyntaxError: invalid syntax"
|
|
mock_container_manager.get_container_logs.return_value = (
|
|
json.dumps({"result": None, "error": error_msg}),
|
|
""
|
|
)
|
|
mock_container_manager.wait_for_container.return_value = 1
|
|
|
|
result = executor.execute("def foo( :")
|
|
|
|
assert result.success is False
|
|
assert "SyntaxError" in result.error
|
|
|
|
|
|
def test_execute_code_with_runtime_error_returns_clear_error(executor, mock_container_manager):
|
|
"""Test code with runtime error returns clear error with traceback."""
|
|
error_msg = "NameError: name 'undefined_var' is not defined"
|
|
mock_container_manager.get_container_logs.return_value = (
|
|
json.dumps({"result": None, "error": error_msg}),
|
|
""
|
|
)
|
|
mock_container_manager.wait_for_container.return_value = 1
|
|
|
|
result = executor.execute("print(undefined_var)")
|
|
|
|
assert result.success is False
|
|
assert "NameError" in result.error
|
|
|
|
|
|
def test_result_serialization_json_compatible_types(executor, mock_container_manager):
|
|
"""Test execution result contains only JSON-serializable data."""
|
|
mock_container_manager.get_container_logs.return_value = (
|
|
json.dumps({"result": [1, 2, {"key": "value"}], "error": None}),
|
|
""
|
|
)
|
|
|
|
result = executor.execute('[1, 2, {"key": "value"}]')
|
|
|
|
# Verify result can be serialized to JSON
|
|
result_dict = result.to_dict()
|
|
json_str = json.dumps(result_dict)
|
|
assert json_str is not None
|
|
|
|
# Verify result data
|
|
assert result.result == [1, 2, {"key": "value"}]
|
|
|
|
|
|
def test_large_output_handling(executor, mock_container_manager):
|
|
"""Test execution handles large output without issues."""
|
|
large_output = "x" * 10000 # 10KB of output
|
|
mock_container_manager.get_container_logs.return_value = (
|
|
json.dumps({"result": None, "error": None}) + "\n" + large_output,
|
|
""
|
|
)
|
|
|
|
result = executor.execute('print("x" * 10000)')
|
|
|
|
assert result.success is True
|
|
assert len(result.stdout) >= 10000
|
|
|
|
|
|
def test_execution_result_to_dict(resource_limits):
|
|
"""Test ExecutionResult.to_dict() returns proper dictionary."""
|
|
result = ExecutionResult(
|
|
success=True,
|
|
stdout="output",
|
|
stderr="",
|
|
result=42,
|
|
execution_time=0.5,
|
|
exit_code=0,
|
|
error=None
|
|
)
|
|
|
|
result_dict = result.to_dict()
|
|
|
|
assert isinstance(result_dict, dict)
|
|
assert result_dict["success"] is True
|
|
assert result_dict["stdout"] == "output"
|
|
assert result_dict["stderr"] == ""
|
|
assert result_dict["result"] == 42
|
|
assert result_dict["execution_time"] == 0.5
|
|
assert result_dict["exit_code"] == 0
|
|
assert result_dict["error"] is None
|
|
|
|
|
|
def test_prepare_code_wraps_code_properly(executor):
|
|
"""Test _prepare_code wraps code to capture result."""
|
|
code = "x = 2 + 2\nx"
|
|
wrapped = executor._prepare_code(code)
|
|
|
|
# Wrapped code should be executable Python
|
|
assert "import" in wrapped
|
|
assert "json" in wrapped
|
|
assert code in wrapped or "2 + 2" in wrapped
|
|
|
|
|
|
def test_parse_output_extracts_result_and_error(executor):
|
|
"""Test _parse_output correctly extracts result and error from JSON."""
|
|
# Test successful result
|
|
stdout = json.dumps({"result": 42, "error": None})
|
|
result, error = executor._parse_output(stdout)
|
|
assert result == 42
|
|
assert error is None
|
|
|
|
# Test error
|
|
stdout = json.dumps({"result": None, "error": "ValueError: invalid"})
|
|
result, error = executor._parse_output(stdout)
|
|
assert result is None
|
|
assert error == "ValueError: invalid"
|
|
|
|
|
|
def test_cleanup_happens_even_on_create_failure(executor, mock_container_manager):
|
|
"""Test container cleanup happens even if create fails."""
|
|
mock_container_manager.create_container.side_effect = Exception("Create failed")
|
|
|
|
with pytest.raises(Exception, match="Create failed"):
|
|
executor.execute("print('test')")
|
|
|
|
# No container to remove since create failed
|
|
mock_container_manager.remove_container.assert_not_called()
|
|
|
|
|
|
def test_cleanup_happens_even_on_start_failure(executor, mock_container_manager):
|
|
"""Test container cleanup happens even if start fails."""
|
|
mock_container_manager.start_container.side_effect = Exception("Start failed")
|
|
|
|
with pytest.raises(Exception, match="Start failed"):
|
|
executor.execute("print('test')")
|
|
|
|
# Container should still be removed
|
|
mock_container_manager.remove_container.assert_called_once_with("test-container-123")
|
|
|
|
|
|
def test_execution_time_tracking(executor, mock_container_manager):
|
|
"""Test execution time is tracked accurately."""
|
|
mock_container_manager.get_container_logs.return_value = (
|
|
json.dumps({"result": None, "error": None}),
|
|
""
|
|
)
|
|
|
|
result = executor.execute("pass")
|
|
|
|
assert result.execution_time >= 0
|
|
assert isinstance(result.execution_time, float)
|
|
|
|
|
|
def test_execute_with_custom_timeout(executor, mock_container_manager):
|
|
"""Test execute respects custom timeout parameter."""
|
|
mock_container_manager.get_container_logs.return_value = (
|
|
json.dumps({"result": None, "error": None}),
|
|
""
|
|
)
|
|
|
|
executor.execute("pass", timeout=60)
|
|
|
|
# Verify wait was called with custom timeout
|
|
mock_container_manager.wait_for_container.assert_called_with("test-container-123", timeout=60)
|
|
|
|
|
|
def test_execute_uses_default_timeout_from_resource_limits(executor, mock_container_manager):
|
|
"""Test execute uses default timeout from resource limits when not specified."""
|
|
mock_container_manager.get_container_logs.return_value = (
|
|
json.dumps({"result": None, "error": None}),
|
|
""
|
|
)
|
|
|
|
executor.execute("pass") # No timeout specified
|
|
|
|
# Should use resource_limits.timeout (30)
|
|
mock_container_manager.wait_for_container.assert_called_with("test-container-123", timeout=30)
|