395 lines
14 KiB
Python
395 lines
14 KiB
Python
"""Tests for the Simple Backend module."""
|
|
|
|
import pytest
|
|
from unittest.mock import Mock, MagicMock, patch
|
|
from pathlib import Path
|
|
|
|
from mcp_forge.execution.simple.backend import SimpleBackend
|
|
from mcp_forge.execution.simple.executor import ExecutionResult
|
|
from mcp_forge.config.schema import ForgeConfig, ExecutionConfig, ImageConfig
|
|
from mcp_forge.podman.containers import SecureContainerManager
|
|
from mcp_forge.security.audit import AuditLogger
|
|
from mcp_forge.security.resource_limits import ResourceLimits
|
|
|
|
|
|
@pytest.fixture
|
|
def mock_config():
|
|
"""Mock ForgeConfig with execution settings."""
|
|
config = Mock(spec=ForgeConfig)
|
|
|
|
# Execution configuration
|
|
config.execution = Mock(spec=ExecutionConfig)
|
|
config.execution.default_timeout = 300
|
|
config.execution.max_timeout = 1800
|
|
config.execution.default_memory = "512m"
|
|
config.execution.max_memory = "2g"
|
|
config.execution.default_cpu_quota = 50000
|
|
config.execution.max_cpu_quota = 100000
|
|
|
|
# Image configuration
|
|
config.images = Mock(spec=ImageConfig)
|
|
config.images.python_3_11 = "mcp-forge/python:3.11"
|
|
config.images.python_3_12 = "mcp-forge/python:3.12"
|
|
|
|
return config
|
|
|
|
|
|
@pytest.fixture
|
|
def mock_container_manager():
|
|
"""Mock SecureContainerManager."""
|
|
return Mock(spec=SecureContainerManager)
|
|
|
|
|
|
@pytest.fixture
|
|
def mock_audit_logger(tmp_path):
|
|
"""Mock AuditLogger."""
|
|
log_file = tmp_path / "audit.log"
|
|
return Mock(spec=AuditLogger)
|
|
|
|
|
|
@pytest.fixture
|
|
def backend(mock_config, mock_container_manager, mock_audit_logger):
|
|
"""SimpleBackend instance with mocked dependencies."""
|
|
return SimpleBackend(
|
|
config=mock_config,
|
|
container_manager=mock_container_manager,
|
|
audit_logger=mock_audit_logger
|
|
)
|
|
|
|
|
|
def test_execute_without_custom_params_uses_defaults(backend, mock_config):
|
|
"""Test execute uses configuration defaults when no params specified."""
|
|
# Mock executor to return a result
|
|
with patch('mcp_forge.execution.simple.backend.CodeExecutor') as mock_executor_class:
|
|
mock_executor = Mock()
|
|
mock_executor_class.return_value = mock_executor
|
|
mock_executor.execute.return_value = ExecutionResult(
|
|
success=True,
|
|
stdout="",
|
|
stderr="",
|
|
result=42,
|
|
execution_time=0.1,
|
|
exit_code=0
|
|
)
|
|
|
|
result = backend.execute("2 + 2")
|
|
|
|
# Verify executor was created with default limits
|
|
mock_executor_class.assert_called_once()
|
|
args, kwargs = mock_executor_class.call_args
|
|
|
|
# Check resource limits
|
|
resource_limits = kwargs.get('resource_limits')
|
|
assert resource_limits is not None
|
|
assert resource_limits.memory_bytes == 512 * 1024 * 1024 # 512m in bytes
|
|
assert resource_limits.cpu_quota == 50000
|
|
assert resource_limits.timeout == 300
|
|
|
|
# Check image
|
|
assert kwargs.get('image') == "mcp-forge/python:3.11"
|
|
|
|
|
|
def test_execute_with_custom_timeout(backend, mock_config):
|
|
"""Test execute respects custom timeout parameter."""
|
|
with patch('mcp_forge.execution.simple.backend.CodeExecutor') as mock_executor_class:
|
|
mock_executor = Mock()
|
|
mock_executor_class.return_value = mock_executor
|
|
mock_executor.execute.return_value = ExecutionResult(
|
|
success=True,
|
|
stdout="",
|
|
stderr="",
|
|
result=None,
|
|
execution_time=0.1,
|
|
exit_code=0
|
|
)
|
|
|
|
backend.execute("pass", timeout=600)
|
|
|
|
# Verify resource limits include custom timeout
|
|
args, kwargs = mock_executor_class.call_args
|
|
resource_limits = kwargs.get('resource_limits')
|
|
assert resource_limits.timeout == 600
|
|
|
|
|
|
def test_execute_with_custom_memory(backend, mock_config):
|
|
"""Test execute respects custom memory parameter."""
|
|
with patch('mcp_forge.execution.simple.backend.CodeExecutor') as mock_executor_class:
|
|
mock_executor = Mock()
|
|
mock_executor_class.return_value = mock_executor
|
|
mock_executor.execute.return_value = ExecutionResult(
|
|
success=True,
|
|
stdout="",
|
|
stderr="",
|
|
result=None,
|
|
execution_time=0.1,
|
|
exit_code=0
|
|
)
|
|
|
|
backend.execute("pass", memory="1g")
|
|
|
|
# Verify resource limits include custom memory
|
|
args, kwargs = mock_executor_class.call_args
|
|
resource_limits = kwargs.get('resource_limits')
|
|
assert resource_limits.memory_bytes == 1024 * 1024 * 1024 # 1g in bytes
|
|
|
|
|
|
def test_execute_with_custom_cpu_quota(backend, mock_config):
|
|
"""Test execute respects custom CPU quota parameter."""
|
|
with patch('mcp_forge.execution.simple.backend.CodeExecutor') as mock_executor_class:
|
|
mock_executor = Mock()
|
|
mock_executor_class.return_value = mock_executor
|
|
mock_executor.execute.return_value = ExecutionResult(
|
|
success=True,
|
|
stdout="",
|
|
stderr="",
|
|
result=None,
|
|
execution_time=0.1,
|
|
exit_code=0
|
|
)
|
|
|
|
backend.execute("pass", cpu_quota=75000)
|
|
|
|
# Verify resource limits include custom CPU quota
|
|
args, kwargs = mock_executor_class.call_args
|
|
resource_limits = kwargs.get('resource_limits')
|
|
assert resource_limits.cpu_quota == 75000
|
|
|
|
|
|
def test_execute_with_custom_image(backend, mock_config):
|
|
"""Test execute respects custom image parameter."""
|
|
with patch('mcp_forge.execution.simple.backend.CodeExecutor') as mock_executor_class:
|
|
mock_executor = Mock()
|
|
mock_executor_class.return_value = mock_executor
|
|
mock_executor.execute.return_value = ExecutionResult(
|
|
success=True,
|
|
stdout="",
|
|
stderr="",
|
|
result=None,
|
|
execution_time=0.1,
|
|
exit_code=0
|
|
)
|
|
|
|
backend.execute("pass", custom_image="mcp-forge/python:3.12")
|
|
|
|
# Verify correct image was used
|
|
args, kwargs = mock_executor_class.call_args
|
|
assert kwargs.get('image') == "mcp-forge/python:3.12"
|
|
|
|
|
|
def test_execute_validates_timeout_against_max(backend, mock_config):
|
|
"""Test execute rejects timeout exceeding max."""
|
|
with pytest.raises(ValueError, match="timeout.*exceeds maximum"):
|
|
backend.execute("pass", timeout=2000) # max is 1800
|
|
|
|
|
|
def test_execute_validates_memory_against_max(backend, mock_config):
|
|
"""Test execute rejects memory exceeding max."""
|
|
with pytest.raises(ValueError, match="memory.*exceeds maximum"):
|
|
backend.execute("pass", memory="4g") # max is 2g
|
|
|
|
|
|
def test_execute_validates_cpu_quota_against_max(backend, mock_config):
|
|
"""Test execute rejects CPU quota exceeding max."""
|
|
with pytest.raises(ValueError, match="cpu_quota.*exceeds maximum"):
|
|
backend.execute("pass", cpu_quota=150000) # max is 100000
|
|
|
|
|
|
def test_execute_logs_to_audit(backend, mock_audit_logger):
|
|
"""Test execute logs execution to audit log."""
|
|
with patch('mcp_forge.execution.simple.backend.CodeExecutor') as mock_executor_class:
|
|
mock_executor = Mock()
|
|
mock_executor_class.return_value = mock_executor
|
|
mock_executor.execute.return_value = ExecutionResult(
|
|
success=True,
|
|
stdout="",
|
|
stderr="",
|
|
result=42,
|
|
execution_time=0.1,
|
|
exit_code=0
|
|
)
|
|
|
|
backend.execute("x = 2 + 2")
|
|
|
|
# Verify audit log was called
|
|
mock_audit_logger.log.assert_called()
|
|
call_args = mock_audit_logger.log.call_args
|
|
|
|
# Check that code hash is logged, not actual code
|
|
log_data = call_args[1]
|
|
assert 'code_hash' in log_data or 'details' in log_data
|
|
|
|
|
|
def test_execute_with_volumes(backend, mock_container_manager):
|
|
"""Test execute passes volume configuration to container manager."""
|
|
with patch('mcp_forge.execution.simple.backend.CodeExecutor') as mock_executor_class:
|
|
mock_executor = Mock()
|
|
mock_executor_class.return_value = mock_executor
|
|
mock_executor.execute.return_value = ExecutionResult(
|
|
success=True,
|
|
stdout="",
|
|
stderr="",
|
|
result=None,
|
|
execution_time=0.1,
|
|
exit_code=0
|
|
)
|
|
|
|
volumes = {
|
|
"/mcp-forge/sessions/test-session/workspace": {"bind": "/workspace", "mode": "rw"}
|
|
}
|
|
|
|
backend.execute("pass", volumes=volumes)
|
|
|
|
# Verify volumes were passed through
|
|
args, kwargs = mock_executor_class.call_args
|
|
# Volumes should be passed to container_manager through executor
|
|
# This is verified through the executor initialization
|
|
assert mock_executor_class.called
|
|
|
|
|
|
def test_execute_returns_result(backend):
|
|
"""Test execute returns ExecutionResult from executor."""
|
|
expected_result = ExecutionResult(
|
|
success=True,
|
|
stdout="Hello\n",
|
|
stderr="",
|
|
result=42,
|
|
execution_time=0.5,
|
|
exit_code=0
|
|
)
|
|
|
|
with patch('mcp_forge.execution.simple.backend.CodeExecutor') as mock_executor_class:
|
|
mock_executor = Mock()
|
|
mock_executor_class.return_value = mock_executor
|
|
mock_executor.execute.return_value = expected_result
|
|
|
|
result = backend.execute('print("Hello"); 42')
|
|
|
|
assert result == expected_result
|
|
assert result.success is True
|
|
assert result.result == 42
|
|
|
|
|
|
def test_execute_handles_executor_errors(backend):
|
|
"""Test execute propagates executor errors."""
|
|
with patch('mcp_forge.execution.simple.backend.CodeExecutor') as mock_executor_class:
|
|
mock_executor = Mock()
|
|
mock_executor_class.return_value = mock_executor
|
|
mock_executor.execute.side_effect = RuntimeError("Container failed")
|
|
|
|
with pytest.raises(RuntimeError, match="Container failed"):
|
|
backend.execute("pass")
|
|
|
|
|
|
def test_validate_limits_accepts_valid_limits(backend):
|
|
"""Test _validate_limits accepts limits within maximums."""
|
|
# Should not raise
|
|
backend._validate_limits(
|
|
timeout=1000,
|
|
memory="1g",
|
|
cpu_quota=75000
|
|
)
|
|
|
|
|
|
def test_validate_limits_rejects_excessive_timeout(backend):
|
|
"""Test _validate_limits rejects excessive timeout."""
|
|
with pytest.raises(ValueError, match="timeout"):
|
|
backend._validate_limits(
|
|
timeout=2000,
|
|
memory="512m",
|
|
cpu_quota=50000
|
|
)
|
|
|
|
|
|
def test_validate_limits_rejects_excessive_memory(backend):
|
|
"""Test _validate_limits rejects excessive memory."""
|
|
with pytest.raises(ValueError, match="memory"):
|
|
backend._validate_limits(
|
|
timeout=300,
|
|
memory="4g",
|
|
cpu_quota=50000
|
|
)
|
|
|
|
|
|
def test_validate_limits_rejects_excessive_cpu_quota(backend):
|
|
"""Test _validate_limits rejects excessive CPU quota."""
|
|
with pytest.raises(ValueError, match="cpu_quota"):
|
|
backend._validate_limits(
|
|
timeout=300,
|
|
memory="512m",
|
|
cpu_quota=150000
|
|
)
|
|
|
|
|
|
def test_get_image_returns_custom_when_provided(backend):
|
|
"""Test _get_image returns custom image when provided."""
|
|
image = backend._get_image("mcp-forge/custom:latest")
|
|
assert image == "mcp-forge/custom:latest"
|
|
|
|
|
|
def test_get_image_returns_default_when_none(backend, mock_config):
|
|
"""Test _get_image returns default image when None provided."""
|
|
image = backend._get_image(None)
|
|
assert image == mock_config.images.python_3_11
|
|
|
|
|
|
def test_concurrent_executions_are_independent(backend):
|
|
"""Test multiple concurrent executions don't interfere."""
|
|
with patch('mcp_forge.execution.simple.backend.CodeExecutor') as mock_executor_class:
|
|
# Create separate mock executors for each call
|
|
executor1 = Mock()
|
|
executor2 = Mock()
|
|
mock_executor_class.side_effect = [executor1, executor2]
|
|
|
|
executor1.execute.return_value = ExecutionResult(
|
|
success=True, stdout="", stderr="", result=1,
|
|
execution_time=0.1, exit_code=0
|
|
)
|
|
executor2.execute.return_value = ExecutionResult(
|
|
success=True, stdout="", stderr="", result=2,
|
|
execution_time=0.1, exit_code=0
|
|
)
|
|
|
|
result1 = backend.execute("1")
|
|
result2 = backend.execute("2")
|
|
|
|
assert result1.result == 1
|
|
assert result2.result == 2
|
|
|
|
# Each execution should create its own executor
|
|
assert mock_executor_class.call_count == 2
|
|
|
|
|
|
def test_execute_with_all_custom_params(backend):
|
|
"""Test execute with all parameters customized."""
|
|
with patch('mcp_forge.execution.simple.backend.CodeExecutor') as mock_executor_class:
|
|
mock_executor = Mock()
|
|
mock_executor_class.return_value = mock_executor
|
|
mock_executor.execute.return_value = ExecutionResult(
|
|
success=True,
|
|
stdout="",
|
|
stderr="",
|
|
result=None,
|
|
execution_time=0.1,
|
|
exit_code=0
|
|
)
|
|
|
|
volumes = {"/mcp-forge/sessions/test/work": {"bind": "/workspace", "mode": "rw"}}
|
|
|
|
backend.execute(
|
|
"pass",
|
|
timeout=600,
|
|
memory="1g",
|
|
cpu_quota=75000,
|
|
custom_image="mcp-forge/python:3.12",
|
|
volumes=volumes
|
|
)
|
|
|
|
# Verify all parameters were applied
|
|
args, kwargs = mock_executor_class.call_args
|
|
|
|
resource_limits = kwargs.get('resource_limits')
|
|
assert resource_limits.timeout == 600
|
|
assert resource_limits.memory_bytes == 1024 * 1024 * 1024 # 1g in bytes
|
|
assert resource_limits.cpu_quota == 75000
|
|
|
|
assert kwargs.get('image') == "mcp-forge/python:3.12"
|