371 lines
13 KiB
Python
371 lines
13 KiB
Python
"""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)
|