325 lines
11 KiB
Python
325 lines
11 KiB
Python
"""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
|