422 lines
14 KiB
Python
422 lines
14 KiB
Python
"""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 pod_executor.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 pod_executor.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
|