mcp-forge/tests/execution/simple/test_executor.py

"""Tests for the Code Executor module."""

import pytest
from unittest.mock import Mock
import json

from pod_executor.simple.executor import CodeExecutor, ExecutionResult
from mcp_forge.podman.containers import SecureContainerManager
from pod_executor.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)