initial commit after one day coding agent session
This commit is contained in:
commit
372af75b90
88 changed files with 22694 additions and 0 deletions
0
tests/execution/jupyter/__init__.py
Normal file
0
tests/execution/jupyter/__init__.py
Normal file
422
tests/execution/jupyter/test_backend.py
Normal file
422
tests/execution/jupyter/test_backend.py
Normal file
|
|
@ -0,0 +1,422 @@
|
|||
"""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
|
||||
325
tests/execution/jupyter/test_kernel.py
Normal file
325
tests/execution/jupyter/test_kernel.py
Normal file
|
|
@ -0,0 +1,325 @@
|
|||
"""Tests for the Jupyter Kernel Manager module."""
|
||||
|
||||
import pytest
|
||||
from unittest.mock import Mock, MagicMock, patch
|
||||
from datetime import datetime, timedelta
|
||||
from pathlib import Path
|
||||
|
||||
from mcp_forge.execution.jupyter.kernel import (
|
||||
JupyterKernelManager,
|
||||
KernelInfo,
|
||||
KernelError
|
||||
)
|
||||
from mcp_forge.podman.containers import SecureContainerManager
|
||||
from mcp_forge.security.resource_limits import ResourceLimits
|
||||
from mcp_forge.execution.simple.executor import ExecutionResult
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def resource_limits():
|
||||
"""Standard resource limits for testing."""
|
||||
return ResourceLimits(
|
||||
memory="512m",
|
||||
cpu_quota=50000,
|
||||
storage="1g",
|
||||
timeout=300
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_container_manager():
|
||||
"""Mock SecureContainerManager."""
|
||||
manager = Mock(spec=SecureContainerManager)
|
||||
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.get_container_logs.return_value = ("", "")
|
||||
return manager
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def kernel_manager(mock_container_manager, resource_limits):
|
||||
"""JupyterKernelManager instance with mocked dependencies."""
|
||||
return JupyterKernelManager(
|
||||
container_manager=mock_container_manager,
|
||||
image="mcp-forge/jupyter:latest",
|
||||
resource_limits=resource_limits
|
||||
)
|
||||
|
||||
|
||||
def test_start_kernel_creates_container(kernel_manager, mock_container_manager):
|
||||
"""Test start_kernel creates and starts a container."""
|
||||
kernel_id = kernel_manager.start_kernel("session-1")
|
||||
|
||||
assert kernel_id is not None
|
||||
assert kernel_id.startswith("kernel-")
|
||||
|
||||
# Verify container was created and started
|
||||
mock_container_manager.create_container.assert_called_once()
|
||||
mock_container_manager.start_container.assert_called_once()
|
||||
|
||||
|
||||
def test_start_kernel_returns_kernel_info(kernel_manager):
|
||||
"""Test start_kernel returns valid kernel info."""
|
||||
kernel_id = kernel_manager.start_kernel("session-1")
|
||||
|
||||
# Kernel should be registered
|
||||
assert kernel_id in kernel_manager.kernels
|
||||
|
||||
kernel_info = kernel_manager.kernels[kernel_id]
|
||||
assert kernel_info.kernel_id == kernel_id
|
||||
assert kernel_info.container_id == "test-container-123"
|
||||
assert kernel_info.session_id == "session-1"
|
||||
assert isinstance(kernel_info.started_at, datetime)
|
||||
|
||||
|
||||
def test_execute_code_in_kernel_returns_result(kernel_manager):
|
||||
"""Test execute_code runs code and returns result."""
|
||||
kernel_id = kernel_manager.start_kernel("session-1")
|
||||
|
||||
result = kernel_manager.execute_code(kernel_id, "2 + 2")
|
||||
|
||||
assert isinstance(result, ExecutionResult)
|
||||
assert result.success is True
|
||||
|
||||
|
||||
def test_execute_code_with_nonexistent_kernel_raises_error(kernel_manager):
|
||||
"""Test execute_code raises error for nonexistent kernel."""
|
||||
with pytest.raises(KernelError, match="Kernel.*not found"):
|
||||
kernel_manager.execute_code("nonexistent-kernel", "pass")
|
||||
|
||||
|
||||
def test_execute_code_preserves_namespace(kernel_manager):
|
||||
"""Test namespace persists between executions."""
|
||||
kernel_id = kernel_manager.start_kernel("session-1")
|
||||
|
||||
# Set a variable
|
||||
result1 = kernel_manager.execute_code(kernel_id, "x = 42")
|
||||
assert result1.success is True
|
||||
|
||||
# Access the variable
|
||||
result2 = kernel_manager.execute_code(kernel_id, "x")
|
||||
assert result2.success is True
|
||||
# In real implementation, result2.result would be 42
|
||||
|
||||
|
||||
def test_shutdown_kernel_removes_container(kernel_manager, mock_container_manager):
|
||||
"""Test shutdown_kernel cleans up container."""
|
||||
kernel_id = kernel_manager.start_kernel("session-1")
|
||||
|
||||
kernel_manager.shutdown_kernel(kernel_id)
|
||||
|
||||
# Verify container was stopped and removed
|
||||
mock_container_manager.stop_container.assert_called_once_with("test-container-123", timeout=10)
|
||||
mock_container_manager.remove_container.assert_called_once_with("test-container-123")
|
||||
|
||||
# Kernel should be removed from registry
|
||||
assert kernel_id not in kernel_manager.kernels
|
||||
|
||||
|
||||
def test_shutdown_nonexistent_kernel_raises_error(kernel_manager):
|
||||
"""Test shutdown_kernel raises error for nonexistent kernel."""
|
||||
with pytest.raises(KernelError, match="Kernel.*not found"):
|
||||
kernel_manager.shutdown_kernel("nonexistent-kernel")
|
||||
|
||||
|
||||
def test_inspect_namespace_returns_variables(kernel_manager):
|
||||
"""Test inspect_namespace returns list of variables."""
|
||||
kernel_id = kernel_manager.start_kernel("session-1")
|
||||
|
||||
# Execute some code to create variables
|
||||
kernel_manager.execute_code(kernel_id, "x = 1; y = 2; z = 3")
|
||||
|
||||
variables = kernel_manager.inspect_namespace(kernel_id)
|
||||
|
||||
assert isinstance(variables, list)
|
||||
# In real implementation, would contain ['x', 'y', 'z']
|
||||
|
||||
|
||||
def test_inspect_namespace_filters_private_vars(kernel_manager):
|
||||
"""Test inspect_namespace filters out private variables."""
|
||||
kernel_id = kernel_manager.start_kernel("session-1")
|
||||
|
||||
kernel_manager.execute_code(kernel_id, "x = 1; _private = 2; __dunder__ = 3")
|
||||
|
||||
variables = kernel_manager.inspect_namespace(kernel_id)
|
||||
|
||||
# Private variables should be filtered
|
||||
# In real implementation: assert '_private' not in variables
|
||||
|
||||
|
||||
def test_get_variable_info_returns_metadata(kernel_manager):
|
||||
"""Test get_variable_info returns variable metadata."""
|
||||
kernel_id = kernel_manager.start_kernel("session-1")
|
||||
|
||||
kernel_manager.execute_code(kernel_id, "x = [1, 2, 3, 4, 5]")
|
||||
|
||||
info = kernel_manager.get_variable_info(kernel_id, "x")
|
||||
|
||||
assert isinstance(info, dict)
|
||||
assert "type" in info
|
||||
# In real implementation: assert info["type"] == "list"
|
||||
|
||||
|
||||
def test_restart_kernel_resets_namespace(kernel_manager, mock_container_manager):
|
||||
"""Test restart_kernel resets namespace but keeps container."""
|
||||
kernel_id = kernel_manager.start_kernel("session-1")
|
||||
original_container_id = kernel_manager.kernels[kernel_id].container_id
|
||||
|
||||
# Set a variable
|
||||
kernel_manager.execute_code(kernel_id, "x = 42")
|
||||
|
||||
# Restart
|
||||
kernel_manager.restart_kernel(kernel_id)
|
||||
|
||||
# Container should be the same
|
||||
assert kernel_manager.kernels[kernel_id].container_id == original_container_id
|
||||
|
||||
# Namespace should be reset (variable no longer accessible)
|
||||
# In real implementation, executing "x" would raise NameError
|
||||
|
||||
|
||||
def test_cleanup_idle_kernels_removes_old_kernels(kernel_manager, mock_container_manager):
|
||||
"""Test cleanup_idle_kernels removes kernels idle too long."""
|
||||
# Start two kernels
|
||||
kernel1 = kernel_manager.start_kernel("session-1")
|
||||
kernel2 = kernel_manager.start_kernel("session-2")
|
||||
|
||||
# Make kernel1 appear old
|
||||
kernel_manager.kernels[kernel1].last_activity = datetime.utcnow() - timedelta(hours=2)
|
||||
|
||||
# Cleanup kernels idle > 1 hour
|
||||
count = kernel_manager.cleanup_idle_kernels(timedelta(hours=1))
|
||||
|
||||
assert count == 1
|
||||
assert kernel1 not in kernel_manager.kernels
|
||||
assert kernel2 in kernel_manager.kernels
|
||||
|
||||
|
||||
def test_kernel_with_volumes(kernel_manager, mock_container_manager):
|
||||
"""Test kernel can be started with volume mounts."""
|
||||
volumes = {
|
||||
"/mcp-forge/sessions/session-1/workspace": {"bind": "/workspace", "mode": "rw"}
|
||||
}
|
||||
|
||||
kernel_id = kernel_manager.start_kernel("session-1", volumes=volumes)
|
||||
|
||||
assert kernel_id is not None
|
||||
# Verify volumes were passed to container creation
|
||||
call_args = mock_container_manager.create_container.call_args
|
||||
# In real implementation, would verify volumes in ContainerConfig
|
||||
|
||||
|
||||
def test_execute_code_with_timeout(kernel_manager):
|
||||
"""Test execute_code respects timeout parameter."""
|
||||
kernel_id = kernel_manager.start_kernel("session-1")
|
||||
|
||||
# Execute with custom timeout
|
||||
result = kernel_manager.execute_code(kernel_id, "import time; time.sleep(0.1)", timeout=10)
|
||||
|
||||
assert isinstance(result, ExecutionResult)
|
||||
|
||||
|
||||
def test_execute_code_handles_syntax_error(kernel_manager):
|
||||
"""Test execute_code handles syntax errors gracefully."""
|
||||
kernel_id = kernel_manager.start_kernel("session-1")
|
||||
|
||||
result = kernel_manager.execute_code(kernel_id, "def foo( :")
|
||||
|
||||
assert result.success is False
|
||||
assert result.error is not None
|
||||
assert "SyntaxError" in result.error or "syntax" in result.error.lower()
|
||||
|
||||
|
||||
def test_execute_code_handles_runtime_error(kernel_manager):
|
||||
"""Test execute_code handles runtime errors gracefully."""
|
||||
kernel_id = kernel_manager.start_kernel("session-1")
|
||||
|
||||
result = kernel_manager.execute_code(kernel_id, "1 / 0")
|
||||
|
||||
assert result.success is False
|
||||
assert result.error is not None
|
||||
|
||||
|
||||
def test_execute_code_captures_stdout(kernel_manager):
|
||||
"""Test execute_code captures stdout output."""
|
||||
kernel_id = kernel_manager.start_kernel("session-1")
|
||||
|
||||
result = kernel_manager.execute_code(kernel_id, 'print("Hello, World!")')
|
||||
|
||||
assert result.success is True
|
||||
# In real implementation: assert "Hello, World!" in result.stdout
|
||||
|
||||
|
||||
def test_execute_code_captures_stderr(kernel_manager):
|
||||
"""Test execute_code captures stderr output."""
|
||||
kernel_id = kernel_manager.start_kernel("session-1")
|
||||
|
||||
result = kernel_manager.execute_code(kernel_id, 'import sys; print("warning", file=sys.stderr)')
|
||||
|
||||
assert result.success is True
|
||||
# In real implementation: assert "warning" in result.stderr
|
||||
|
||||
|
||||
def test_multiple_kernels_are_isolated(kernel_manager):
|
||||
"""Test multiple kernels have isolated namespaces."""
|
||||
kernel1 = kernel_manager.start_kernel("session-1")
|
||||
kernel2 = kernel_manager.start_kernel("session-2")
|
||||
|
||||
# Set variable in kernel1
|
||||
kernel_manager.execute_code(kernel1, "x = 1")
|
||||
|
||||
# Set different value in kernel2
|
||||
kernel_manager.execute_code(kernel2, "x = 2")
|
||||
|
||||
# Values should be independent
|
||||
result1 = kernel_manager.execute_code(kernel1, "x")
|
||||
result2 = kernel_manager.execute_code(kernel2, "x")
|
||||
|
||||
# In real implementation: verify result1.result == 1 and result2.result == 2
|
||||
|
||||
|
||||
def test_kernel_info_to_dict(resource_limits):
|
||||
"""Test KernelInfo.to_dict() serialization."""
|
||||
now = datetime.utcnow()
|
||||
kernel_info = KernelInfo(
|
||||
kernel_id="kernel-123",
|
||||
container_id="container-456",
|
||||
session_id="session-789",
|
||||
started_at=now,
|
||||
last_activity=now
|
||||
)
|
||||
|
||||
info_dict = kernel_info.to_dict()
|
||||
|
||||
assert isinstance(info_dict, dict)
|
||||
assert info_dict["kernel_id"] == "kernel-123"
|
||||
assert info_dict["container_id"] == "container-456"
|
||||
assert info_dict["session_id"] == "session-789"
|
||||
assert "started_at" in info_dict
|
||||
assert "last_activity" in info_dict
|
||||
|
||||
|
||||
def test_start_kernel_with_session_id_tracking(kernel_manager):
|
||||
"""Test kernel tracks session_id correctly."""
|
||||
kernel_id = kernel_manager.start_kernel("my-session")
|
||||
|
||||
kernel_info = kernel_manager.kernels[kernel_id]
|
||||
assert kernel_info.session_id == "my-session"
|
||||
|
||||
|
||||
def test_update_activity_timestamp(kernel_manager):
|
||||
"""Test executing code updates last_activity timestamp."""
|
||||
kernel_id = kernel_manager.start_kernel("session-1")
|
||||
|
||||
original_activity = kernel_manager.kernels[kernel_id].last_activity
|
||||
|
||||
# Small delay to ensure timestamp difference
|
||||
import time
|
||||
time.sleep(0.01)
|
||||
|
||||
kernel_manager.execute_code(kernel_id, "pass")
|
||||
|
||||
new_activity = kernel_manager.kernels[kernel_id].last_activity
|
||||
assert new_activity > original_activity
|
||||
371
tests/execution/jupyter/test_sessions.py
Normal file
371
tests/execution/jupyter/test_sessions.py
Normal file
|
|
@ -0,0 +1,371 @@
|
|||
"""Tests for the Session Manager module."""
|
||||
|
||||
import pytest
|
||||
from unittest.mock import Mock, MagicMock
|
||||
from datetime import datetime, timedelta
|
||||
from pathlib import Path
|
||||
|
||||
from mcp_forge.execution.jupyter.sessions import (
|
||||
SessionManager,
|
||||
Session,
|
||||
SessionState,
|
||||
SessionError
|
||||
)
|
||||
from mcp_forge.execution.jupyter.kernel import JupyterKernelManager
|
||||
from mcp_forge.config.schema import SessionConfig
|
||||
from mcp_forge.security.audit import AuditLogger
|
||||
from mcp_forge.security.resource_limits import ResourceLimits
|
||||
from mcp_forge.execution.simple.executor import ExecutionResult
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def session_config():
|
||||
"""Mock SessionConfig."""
|
||||
config = Mock(spec=SessionConfig)
|
||||
config.idle_timeout = 3600
|
||||
config.max_concurrent = 10
|
||||
config.cleanup_interval = 300
|
||||
return config
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_kernel_manager():
|
||||
"""Mock JupyterKernelManager."""
|
||||
manager = Mock(spec=JupyterKernelManager)
|
||||
manager.start_kernel.return_value = "kernel-123"
|
||||
manager.execute_code.return_value = ExecutionResult(
|
||||
success=True,
|
||||
stdout="",
|
||||
stderr="",
|
||||
result=None,
|
||||
execution_time=0.1,
|
||||
exit_code=0
|
||||
)
|
||||
manager.shutdown_kernel.return_value = None
|
||||
manager.inspect_namespace.return_value = []
|
||||
return manager
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_audit_logger(tmp_path):
|
||||
"""Mock AuditLogger."""
|
||||
log_file = tmp_path / "audit.log"
|
||||
return Mock(spec=AuditLogger)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def resource_limits():
|
||||
"""Standard resource limits."""
|
||||
return ResourceLimits(
|
||||
memory="512m",
|
||||
cpu_quota=50000,
|
||||
storage="1g",
|
||||
timeout=300
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def session_manager(session_config, mock_kernel_manager, mock_audit_logger):
|
||||
"""SessionManager instance with mocked dependencies."""
|
||||
return SessionManager(
|
||||
config=session_config,
|
||||
kernel_manager=mock_kernel_manager,
|
||||
audit_logger=mock_audit_logger
|
||||
)
|
||||
|
||||
|
||||
def test_create_session_starts_kernel(session_manager, mock_kernel_manager, resource_limits):
|
||||
"""Test create_session starts a kernel."""
|
||||
session = session_manager.create_session("session-1", resource_limits)
|
||||
|
||||
assert session.session_id == "session-1"
|
||||
assert session.kernel_id == "kernel-123"
|
||||
mock_kernel_manager.start_kernel.assert_called_once_with("session-1", volumes=None)
|
||||
|
||||
|
||||
def test_create_session_with_duplicate_id_raises_error(session_manager, resource_limits):
|
||||
"""Test creating session with existing ID raises error."""
|
||||
session_manager.create_session("session-1", resource_limits)
|
||||
|
||||
with pytest.raises(SessionError, match="already exists"):
|
||||
session_manager.create_session("session-1", resource_limits)
|
||||
|
||||
|
||||
def test_create_session_enforces_max_concurrent(session_manager, session_config, resource_limits):
|
||||
"""Test max concurrent sessions is enforced."""
|
||||
session_config.max_concurrent = 2
|
||||
|
||||
# Create 2 sessions (at limit)
|
||||
session_manager.create_session("session-1", resource_limits)
|
||||
session_manager.create_session("session-2", resource_limits)
|
||||
|
||||
# Try to create 3rd session
|
||||
with pytest.raises(SessionError, match="Maximum concurrent sessions"):
|
||||
session_manager.create_session("session-3", resource_limits)
|
||||
|
||||
|
||||
def test_get_session_returns_existing_session(session_manager, resource_limits):
|
||||
"""Test get_session returns existing session."""
|
||||
created = session_manager.create_session("session-1", resource_limits)
|
||||
retrieved = session_manager.get_session("session-1")
|
||||
|
||||
assert retrieved.session_id == created.session_id
|
||||
assert retrieved.kernel_id == created.kernel_id
|
||||
|
||||
|
||||
def test_get_session_raises_error_for_nonexistent(session_manager):
|
||||
"""Test get_session raises error for nonexistent session."""
|
||||
with pytest.raises(SessionError, match="not found"):
|
||||
session_manager.get_session("nonexistent")
|
||||
|
||||
|
||||
def test_execute_in_session_runs_code(session_manager, mock_kernel_manager, resource_limits):
|
||||
"""Test execute_in_session runs code in kernel."""
|
||||
session_manager.create_session("session-1", resource_limits)
|
||||
|
||||
result = session_manager.execute_in_session("session-1", "x = 42")
|
||||
|
||||
assert result.success is True
|
||||
mock_kernel_manager.execute_code.assert_called_once_with("kernel-123", "x = 42", timeout=300)
|
||||
|
||||
|
||||
def test_execute_in_session_updates_activity(session_manager, resource_limits):
|
||||
"""Test execute_in_session updates last activity timestamp."""
|
||||
session_manager.create_session("session-1", resource_limits)
|
||||
session = session_manager.get_session("session-1")
|
||||
|
||||
original_activity = session.last_activity
|
||||
|
||||
import time
|
||||
time.sleep(0.01)
|
||||
|
||||
session_manager.execute_in_session("session-1", "pass")
|
||||
|
||||
assert session.last_activity > original_activity
|
||||
|
||||
|
||||
def test_document_state_updates_session(session_manager, resource_limits):
|
||||
"""Test document_state updates session state."""
|
||||
session_manager.create_session("session-1", resource_limits)
|
||||
|
||||
variables = {"x": "The result of computation", "y": "Another variable"}
|
||||
session_manager.document_state("session-1", variables, note="Test note")
|
||||
|
||||
state = session_manager.get_session_state("session-1")
|
||||
|
||||
assert state.documented_variables == variables
|
||||
assert state.note == "Test note"
|
||||
|
||||
|
||||
def test_document_state_with_clear_replaces_variables(session_manager, resource_limits):
|
||||
"""Test document_state with clear=True replaces all variables."""
|
||||
session_manager.create_session("session-1", resource_limits)
|
||||
|
||||
# Set initial variables
|
||||
session_manager.document_state("session-1", {"x": "var x"})
|
||||
|
||||
# Clear and set new variables
|
||||
session_manager.document_state("session-1", {"y": "var y"}, clear=True)
|
||||
|
||||
state = session_manager.get_session_state("session-1")
|
||||
|
||||
assert "x" not in state.documented_variables
|
||||
assert "y" in state.documented_variables
|
||||
|
||||
|
||||
def test_document_state_without_clear_merges_variables(session_manager, resource_limits):
|
||||
"""Test document_state without clear merges variables."""
|
||||
session_manager.create_session("session-1", resource_limits)
|
||||
|
||||
session_manager.document_state("session-1", {"x": "var x"})
|
||||
session_manager.document_state("session-1", {"y": "var y"})
|
||||
|
||||
state = session_manager.get_session_state("session-1")
|
||||
|
||||
assert "x" in state.documented_variables
|
||||
assert "y" in state.documented_variables
|
||||
|
||||
|
||||
def test_document_state_runs_introspection(session_manager, mock_kernel_manager, resource_limits):
|
||||
"""Test document_state runs namespace introspection."""
|
||||
mock_kernel_manager.inspect_namespace.return_value = ["x", "y", "z"]
|
||||
mock_kernel_manager.get_variable_info.return_value = {"type": "int", "repr": "42"}
|
||||
|
||||
session_manager.create_session("session-1", resource_limits)
|
||||
session_manager.document_state("session-1", {"x": "documented"})
|
||||
|
||||
state = session_manager.get_session_state("session-1")
|
||||
|
||||
assert state.all_variables == ["x", "y", "z"]
|
||||
mock_kernel_manager.inspect_namespace.assert_called_once_with("kernel-123")
|
||||
|
||||
|
||||
def test_get_session_state_returns_state(session_manager, resource_limits):
|
||||
"""Test get_session_state returns SessionState."""
|
||||
session_manager.create_session("session-1", resource_limits)
|
||||
|
||||
state = session_manager.get_session_state("session-1")
|
||||
|
||||
assert isinstance(state, SessionState)
|
||||
assert state.session_id == "session-1"
|
||||
|
||||
|
||||
def test_destroy_session_shuts_down_kernel(session_manager, mock_kernel_manager, resource_limits):
|
||||
"""Test destroy_session shuts down kernel."""
|
||||
session_manager.create_session("session-1", resource_limits)
|
||||
|
||||
session_manager.destroy_session("session-1")
|
||||
|
||||
mock_kernel_manager.shutdown_kernel.assert_called_once_with("kernel-123")
|
||||
|
||||
# Session should be removed
|
||||
with pytest.raises(SessionError):
|
||||
session_manager.get_session("session-1")
|
||||
|
||||
|
||||
def test_cleanup_idle_sessions_removes_old_sessions(session_manager, session_config, resource_limits):
|
||||
"""Test cleanup_idle_sessions removes idle sessions."""
|
||||
session_config.idle_timeout = 3600 # 1 hour
|
||||
|
||||
# Create two sessions
|
||||
session_manager.create_session("session-1", resource_limits)
|
||||
session_manager.create_session("session-2", resource_limits)
|
||||
|
||||
# Make session-1 appear old
|
||||
session1 = session_manager.get_session("session-1")
|
||||
session1.last_activity = datetime.utcnow() - timedelta(hours=2)
|
||||
|
||||
# Cleanup
|
||||
count = session_manager.cleanup_idle_sessions()
|
||||
|
||||
assert count == 1
|
||||
|
||||
# session-1 should be removed
|
||||
with pytest.raises(SessionError):
|
||||
session_manager.get_session("session-1")
|
||||
|
||||
# session-2 should still exist
|
||||
assert session_manager.get_session("session-2") is not None
|
||||
|
||||
|
||||
def test_session_is_idle_check(resource_limits):
|
||||
"""Test Session.is_idle() check."""
|
||||
now = datetime.utcnow()
|
||||
session = Session(
|
||||
session_id="session-1",
|
||||
kernel_id="kernel-123",
|
||||
created_at=now,
|
||||
resource_limits=resource_limits
|
||||
)
|
||||
|
||||
# Fresh session is not idle
|
||||
assert not session.is_idle(timedelta(hours=1))
|
||||
|
||||
# Make it old
|
||||
session.last_activity = now - timedelta(hours=2)
|
||||
|
||||
# Now it's idle
|
||||
assert session.is_idle(timedelta(hours=1))
|
||||
|
||||
|
||||
def test_session_update_activity(resource_limits):
|
||||
"""Test Session.update_activity() updates timestamp."""
|
||||
now = datetime.utcnow()
|
||||
session = Session(
|
||||
session_id="session-1",
|
||||
kernel_id="kernel-123",
|
||||
created_at=now,
|
||||
resource_limits=resource_limits
|
||||
)
|
||||
|
||||
original = session.last_activity
|
||||
|
||||
import time
|
||||
time.sleep(0.01)
|
||||
|
||||
session.update_activity()
|
||||
|
||||
assert session.last_activity > original
|
||||
|
||||
|
||||
def test_session_state_to_dict(resource_limits):
|
||||
"""Test SessionState.to_dict() serialization."""
|
||||
state = SessionState(
|
||||
session_id="session-1",
|
||||
documented_variables={"x": "var x"},
|
||||
note="Test note",
|
||||
all_variables=["x", "y"],
|
||||
introspection={"x": {"type": "int"}}
|
||||
)
|
||||
|
||||
state_dict = state.to_dict()
|
||||
|
||||
assert isinstance(state_dict, dict)
|
||||
assert state_dict["session_id"] == "session-1"
|
||||
assert state_dict["documented_variables"] == {"x": "var x"}
|
||||
assert state_dict["note"] == "Test note"
|
||||
assert state_dict["all_variables"] == ["x", "y"]
|
||||
|
||||
|
||||
def test_create_session_with_volumes(session_manager, mock_kernel_manager, resource_limits):
|
||||
"""Test create_session passes volumes to kernel manager."""
|
||||
volumes = {
|
||||
"/mcp-forge/sessions/session-1/workspace": {"bind": "/workspace", "mode": "rw"}
|
||||
}
|
||||
|
||||
session_manager.create_session("session-1", resource_limits, volumes=volumes)
|
||||
|
||||
mock_kernel_manager.start_kernel.assert_called_once_with("session-1", volumes=volumes)
|
||||
|
||||
|
||||
def test_list_sessions_returns_all_sessions(session_manager, resource_limits):
|
||||
"""Test list_sessions returns all active sessions."""
|
||||
session_manager.create_session("session-1", resource_limits)
|
||||
session_manager.create_session("session-2", resource_limits)
|
||||
|
||||
sessions = session_manager.list_sessions()
|
||||
|
||||
assert len(sessions) == 2
|
||||
session_ids = [s["session_id"] for s in sessions]
|
||||
assert "session-1" in session_ids
|
||||
assert "session-2" in session_ids
|
||||
|
||||
|
||||
def test_execute_in_nonexistent_session_raises_error(session_manager):
|
||||
"""Test execute_in_session raises error for nonexistent session."""
|
||||
with pytest.raises(SessionError, match="not found"):
|
||||
session_manager.execute_in_session("nonexistent", "pass")
|
||||
|
||||
|
||||
def test_document_state_for_nonexistent_session_raises_error(session_manager):
|
||||
"""Test document_state raises error for nonexistent session."""
|
||||
with pytest.raises(SessionError, match="not found"):
|
||||
session_manager.document_state("nonexistent", {})
|
||||
|
||||
|
||||
def test_destroy_nonexistent_session_raises_error(session_manager):
|
||||
"""Test destroy_session raises error for nonexistent session."""
|
||||
with pytest.raises(SessionError, match="not found"):
|
||||
session_manager.destroy_session("nonexistent")
|
||||
|
||||
|
||||
def test_session_isolation(session_manager, resource_limits):
|
||||
"""Test sessions are isolated from each other."""
|
||||
session_manager.create_session("session-1", resource_limits)
|
||||
session_manager.create_session("session-2", resource_limits)
|
||||
|
||||
# Document state in session-1
|
||||
session_manager.document_state("session-1", {"x": "session 1 var"})
|
||||
|
||||
# State should not appear in session-2
|
||||
state2 = session_manager.get_session_state("session-2")
|
||||
assert "x" not in state2.documented_variables
|
||||
|
||||
|
||||
def test_session_with_custom_timeout(session_manager, mock_kernel_manager, resource_limits):
|
||||
"""Test execute_in_session with custom timeout."""
|
||||
session_manager.create_session("session-1", resource_limits)
|
||||
|
||||
session_manager.execute_in_session("session-1", "pass", timeout=600)
|
||||
|
||||
mock_kernel_manager.execute_code.assert_called_once_with("kernel-123", "pass", timeout=600)
|
||||
Loading…
Add table
Add a link
Reference in a new issue