mcp-forge/tests/execution/jupyter/test_backend.py

"""Tests for Jupyter Backend module."""

import pytest
from unittest.mock import Mock, MagicMock, patch
from pathlib import Path

from mcp_forge.execution.jupyter.backend import JupyterBackend
from mcp_forge.execution.jupyter.sessions import SessionManager, Session, SessionState
from mcp_forge.execution.simple.executor import ExecutionResult
from mcp_forge.config.schema import ForgeConfig, ExecutionConfig, ImageConfig, SessionConfig
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.jupyter = "mcp-forge/jupyter:latest"
    
    # Session configuration
    config.sessions = Mock(spec=SessionConfig)
    config.sessions.idle_timeout = 3600
    config.sessions.max_concurrent = 10
    config.sessions.cleanup_interval = 300
    
    return config


@pytest.fixture
def mock_container_manager():
    """Mock SecureContainerManager."""
    return Mock(spec=SecureContainerManager)


@pytest.fixture
def mock_audit_logger():
    """Mock AuditLogger."""
    return Mock(spec=AuditLogger)


@pytest.fixture
def mock_session_manager():
    """Mock SessionManager."""
    return Mock(spec=SessionManager)


@pytest.fixture
def backend(mock_config, mock_container_manager, mock_audit_logger):
    """JupyterBackend instance with mocked dependencies."""
    with patch('mcp_forge.execution.jupyter.backend.SessionManager') as mock_sm_class:
        mock_session_manager = Mock(spec=SessionManager)
        mock_sm_class.return_value = mock_session_manager
        
        backend = JupyterBackend(
            config=mock_config,
            container_manager=mock_container_manager,
            audit_logger=mock_audit_logger
        )
        backend.session_manager = mock_session_manager
        
        return backend


def test_backend_initializes_session_manager(mock_config, mock_container_manager, mock_audit_logger):
    """Test backend creates SessionManager on initialization."""
    with patch('mcp_forge.execution.jupyter.backend.SessionManager') as mock_sm_class:
        with patch('mcp_forge.execution.jupyter.backend.JupyterKernelManager') as mock_km_class:
            mock_session_manager = Mock(spec=SessionManager)
            mock_sm_class.return_value = mock_session_manager
            
            backend = JupyterBackend(
                config=mock_config,
                container_manager=mock_container_manager,
                audit_logger=mock_audit_logger
            )
            
            # Verify kernel manager was created
            mock_km_class.assert_called_once()
            
            # Verify session manager was created
            mock_sm_class.assert_called_once()


def test_execute_creates_session_if_not_exists(backend):
    """Test execute creates new session if it doesn't exist."""
    # Mock get_session to raise SessionError (session doesn't exist)
    from mcp_forge.execution.jupyter.sessions import SessionError
    backend.session_manager.get_session.side_effect = SessionError("Session not found")
    
    # Mock create_session
    mock_session = Mock(spec=Session)
    backend.session_manager.create_session.return_value = mock_session
    
    # Mock execute_in_session
    mock_result = ExecutionResult(
        success=True,
        stdout="Hello",
        stderr="",
        result="Hello",
        execution_time=0.1,
        exit_code=0
    )
    backend.session_manager.execute_in_session.return_value = mock_result
    
    result = backend.execute("print('Hello')", session_id="test-session")
    
    # Verify session was created
    backend.session_manager.create_session.assert_called_once()
    
    # Verify execution happened
    backend.session_manager.execute_in_session.assert_called_once_with(
        session_id="test-session", code="print('Hello')", timeout=300
    )
    
    assert result.success is True


def test_execute_reuses_existing_session(backend):
    """Test execute reuses existing session."""
    # Mock get_session to return existing session
    mock_session = Mock(spec=Session)
    backend.session_manager.get_session.return_value = mock_session
    
    # Mock execute_in_session
    mock_result = ExecutionResult(
        success=True,
        stdout="42",
        stderr="",
        result=42,
        execution_time=0.1,
        exit_code=0
    )
    backend.session_manager.execute_in_session.return_value = mock_result
    
    result = backend.execute("21 + 21", session_id="existing-session")
    
    # Verify session was NOT created
    backend.session_manager.create_session.assert_not_called()
    
    # Verify session was checked
    backend.session_manager.get_session.assert_called_once_with("existing-session")
    
    # Verify execution happened
    backend.session_manager.execute_in_session.assert_called_once()
    
    assert result.result == 42


def test_execute_with_custom_timeout(backend):
    """Test execute respects custom timeout parameter."""
    # Mock existing session
    mock_session = Mock(spec=Session)
    backend.session_manager.get_session.return_value = mock_session
    
    mock_result = ExecutionResult(
        success=True,
        stdout="",
        stderr="",
        result=None,
        execution_time=0.1,
        exit_code=0
    )
    backend.session_manager.execute_in_session.return_value = mock_result
    
    backend.execute("pass", session_id="test", timeout=600)
    
    # Verify timeout was passed through
    backend.session_manager.execute_in_session.assert_called_once_with(
        session_id="test", code="pass", timeout=600
    )


def test_execute_with_custom_memory(backend, mock_config):
    """Test execute creates session with custom memory limit."""
    from mcp_forge.execution.jupyter.sessions import SessionError
    backend.session_manager.get_session.side_effect = SessionError("Not found")
    
    mock_session = Mock(spec=Session)
    backend.session_manager.create_session.return_value = mock_session
    
    mock_result = ExecutionResult(
        success=True,
        stdout="",
        stderr="",
        result=None,
        execution_time=0.1,
        exit_code=0
    )
    backend.session_manager.execute_in_session.return_value = mock_result
    
    backend.execute("pass", session_id="test", memory="1g")
    
    # Verify session was created with custom memory
    backend.session_manager.create_session.assert_called_once()
    call_args = backend.session_manager.create_session.call_args
    resource_limits = call_args[1]['resource_limits']
    assert resource_limits.memory_bytes == 1024 * 1024 * 1024  # 1g in bytes


def test_execute_with_custom_cpu_quota(backend):
    """Test execute creates session with custom CPU quota."""
    from mcp_forge.execution.jupyter.sessions import SessionError
    backend.session_manager.get_session.side_effect = SessionError("Not found")
    
    mock_session = Mock(spec=Session)
    backend.session_manager.create_session.return_value = mock_session
    
    mock_result = ExecutionResult(
        success=True,
        stdout="",
        stderr="",
        result=None,
        execution_time=0.1,
        exit_code=0
    )
    backend.session_manager.execute_in_session.return_value = mock_result
    
    backend.execute("pass", session_id="test", cpu_quota=75000)
    
    # Verify session was created with custom CPU quota
    backend.session_manager.create_session.assert_called_once()
    call_args = backend.session_manager.create_session.call_args
    resource_limits = call_args[1]['resource_limits']
    assert resource_limits.cpu_quota == 75000


def test_execute_validates_timeout_against_max(backend, mock_config):
    """Test execute rejects timeout exceeding maximum."""
    mock_config.execution.max_timeout = 1800
    
    with pytest.raises(ValueError, match="Timeout 3600 exceeds maximum"):
        backend.execute("pass", session_id="test", timeout=3600)


def test_execute_validates_memory_against_max(backend, mock_config):
    """Test execute rejects memory exceeding maximum."""
    mock_config.execution.max_memory = "2g"
    
    with pytest.raises(ValueError, match="Memory 4g exceeds maximum"):
        backend.execute("pass", session_id="test", memory="4g")


def test_execute_validates_cpu_quota_against_max(backend, mock_config):
    """Test execute rejects CPU quota exceeding maximum."""
    mock_config.execution.max_cpu_quota = 100000
    
    with pytest.raises(ValueError, match="CPU quota 150000 exceeds maximum"):
        backend.execute("pass", session_id="test", cpu_quota=150000)


def test_execute_with_volumes(backend):
    """Test execute passes volumes to session creation."""
    from mcp_forge.execution.jupyter.sessions import SessionError
    backend.session_manager.get_session.side_effect = SessionError("Not found")
    
    mock_session = Mock(spec=Session)
    backend.session_manager.create_session.return_value = mock_session
    
    mock_result = ExecutionResult(
        success=True,
        stdout="",
        stderr="",
        result=None,
        execution_time=0.1,
        exit_code=0
    )
    backend.session_manager.execute_in_session.return_value = mock_result
    
    volumes = {"/host/path": {"bind": "/container/path", "mode": "ro"}}
    backend.execute("pass", session_id="test", volumes=volumes)
    
    # Verify volumes were passed to create_session
    backend.session_manager.create_session.assert_called_once()
    call_args = backend.session_manager.create_session.call_args
    assert call_args[1]['volumes'] == volumes


def test_document_state(backend):
    """Test document_state delegates to session manager."""
    variables = {"x": "Input data", "y": "Output result"}
    note = "Initial data load"
    
    backend.document_state("test-session", variables, note=note, clear=False)
    
    backend.session_manager.document_state.assert_called_once_with(
        session_id="test-session", variables=variables, note=note, clear=False
    )


def test_document_state_with_clear(backend):
    """Test document_state with clear flag."""
    variables = {"new_var": "New data"}
    
    backend.document_state("test-session", variables, clear=True)
    
    backend.session_manager.document_state.assert_called_once_with(
        session_id="test-session", variables=variables, note="", clear=True
    )


def test_get_session_state(backend):
    """Test get_session_state delegates to session manager."""
    mock_state = Mock(spec=SessionState)
    backend.session_manager.get_session_state.return_value = mock_state
    
    state = backend.get_session_state("test-session")
    
    backend.session_manager.get_session_state.assert_called_once_with("test-session")
    assert state == mock_state


def test_destroy_session(backend):
    """Test destroy_session delegates to session manager."""
    backend.destroy_session("test-session")
    
    backend.session_manager.destroy_session.assert_called_once_with("test-session")


def test_list_sessions(backend):
    """Test list_sessions delegates to session manager."""
    mock_sessions = [
        {"session_id": "session1", "kernel_id": "kernel1"},
        {"session_id": "session2", "kernel_id": "kernel2"}
    ]
    backend.session_manager.list_sessions.return_value = mock_sessions
    
    sessions = backend.list_sessions()
    
    backend.session_manager.list_sessions.assert_called_once()
    assert sessions == mock_sessions


def test_cleanup_idle_sessions(backend):
    """Test cleanup_idle_sessions delegates to session manager."""
    backend.session_manager.cleanup_idle_sessions.return_value = 2
    
    count = backend.cleanup_idle_sessions()
    
    backend.session_manager.cleanup_idle_sessions.assert_called_once()
    assert count == 2


def test_default_resource_limits_from_config(mock_config, mock_container_manager, mock_audit_logger):
    """Test _default_resource_limits creates limits from config."""
    with patch('mcp_forge.execution.jupyter.backend.SessionManager'):
        backend = JupyterBackend(
            config=mock_config,
            container_manager=mock_container_manager,
            audit_logger=mock_audit_logger
        )
        
        limits = backend._default_resource_limits()
        
        assert limits.memory_bytes == 512 * 1024 * 1024  # 512m
        assert limits.cpu_quota == 50000
        assert limits.timeout == 300


def test_execute_logs_audit_event(backend):
    """Test execute logs audit event."""
    # Mock existing session
    mock_session = Mock(spec=Session)
    backend.session_manager.get_session.return_value = mock_session
    
    mock_result = ExecutionResult(
        success=True,
        stdout="",
        stderr="",
        result=None,
        execution_time=0.1,
        exit_code=0
    )
    backend.session_manager.execute_in_session.return_value = mock_result
    
    backend.execute("print('test')", session_id="test-session")
    
    # Verify audit log was called
    backend.audit_logger.log.assert_called()


def test_execute_with_defaults_uses_config_values(backend, mock_config):
    """Test execute without parameters uses config defaults."""
    from mcp_forge.execution.jupyter.sessions import SessionError
    backend.session_manager.get_session.side_effect = SessionError("Not found")
    
    mock_session = Mock(spec=Session)
    backend.session_manager.create_session.return_value = mock_session
    
    mock_result = ExecutionResult(
        success=True,
        stdout="",
        stderr="",
        result=None,
        execution_time=0.1,
        exit_code=0
    )
    backend.session_manager.execute_in_session.return_value = mock_result
    
    backend.execute("pass", session_id="test")
    
    # Verify default values from config were used
    call_args = backend.session_manager.create_session.call_args
    resource_limits = call_args[1]['resource_limits']
    assert resource_limits.memory_bytes == 512 * 1024 * 1024
    assert resource_limits.cpu_quota == 50000
    
    call_args = backend.session_manager.execute_in_session.call_args
    assert call_args[1]['timeout'] == 300