"""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)