"""Tests for MCP Context implementation.""" import asyncio from unittest.mock import AsyncMock, Mock import pytest from arcade_mcp_server.context import Context from arcade_mcp_server.context import get_current_model_context as get_current_context from arcade_mcp_server.context import set_current_model_context as set_current_context from arcade_mcp_server.types import ( ModelHint, ModelPreferences, ) class TestContext: """Test Context class and context management.""" def test_context_creation(self, mcp_server): """Test context creation with various parameters.""" # Basic context context = Context(server=mcp_server) assert context.server == mcp_server assert context.request_id is None assert context.session_id is None # Context with request ID context2 = Context(server=mcp_server, request_id="req-123") assert context2.request_id == "req-123" assert context2.session_id is None # No session set yet def test_context_implements_protocol(self): """Test that Context implements MCPContext protocol.""" # Context should expose namespaced adapters assert hasattr(Context, "log") assert hasattr(Context, "progress") assert hasattr(Context, "resources") assert hasattr(Context, "tools") assert hasattr(Context, "prompts") assert hasattr(Context, "sampling") assert hasattr(Context, "ui") assert hasattr(Context, "notifications") def test_context_var_management(self): """Test context variable get/set functionality.""" server = Mock() context = Context(server=server) # Initially no current context assert get_current_context() is None # Set context token = set_current_context(context) assert get_current_context() == context # Clear context set_current_context(None, token) assert get_current_context() is None @pytest.mark.asyncio async def test_context_isolation(self): """Test that contexts are isolated between async tasks.""" server = Mock() context1 = Context(server=server, request_id="req-1") context2 = Context(server=server, request_id="req-2") results = [] async def task1(): set_current_context(context1) results.append(get_current_context()) async def task2(): set_current_context(context2) results.append(get_current_context()) # Run tasks await asyncio.gather(task1(), task2()) # Each task should have its own context assert len(results) == 2 # Context vars are task-local, so both should see their own context assert context1 in results assert context2 in results @pytest.mark.asyncio async def test_logging_methods(self, mcp_server): """Test logging methods.""" session = Mock() session.send_log_message = AsyncMock() context = Context(server=mcp_server) context.set_session(session) # Test all log levels await context.log.debug("Debug message") await context.log.info("Info message") await context.log.warning("Warning message") await context.log.error("Error message") # Verify calls assert session.send_log_message.call_count == 4 # Test with extra metadata await context.log("info", "Test message", logger_name="test.logger", extra={"key": "value"}) # Check the call - context passes logger_name but session expects logger call_kwargs = session.send_log_message.call_args[1] assert call_kwargs["level"] == "info" assert isinstance(call_kwargs["data"], dict) assert call_kwargs["data"]["msg"] == "Test message" assert call_kwargs["data"]["extra"] == {"key": "value"} assert call_kwargs["logger"] == "test.logger" @pytest.mark.asyncio async def test_logging_without_session(self, mcp_server): """Test logging when session is not available.""" context = Context(server=mcp_server) # Should not raise errors await context.log.debug("Debug message") await context.log.info("Info message") await context.log.warning("Warning message") await context.log.error("Error message") @pytest.mark.asyncio async def test_progress_reporting(self, mcp_server): """Test progress reporting functionality.""" session = Mock() session.send_progress_notification = AsyncMock() session._request_meta = Mock(progressToken="task-123") context = Context(server=mcp_server) context.set_session(session) # Report progress await context.progress.report(50, 100, "Processing...") session.send_progress_notification.assert_called_once_with( progress_token="task-123", progress=50, total=100, message="Processing..." ) # Without total await context.progress.report(0.75, message="Almost done") assert session.send_progress_notification.call_count == 2 # Test without progress token - should not call send_progress_notification session2 = Mock(spec=["send_progress_notification"]) session2.send_progress_notification = AsyncMock() # Without _request_meta attribute, progress won't be reported context2 = Context(server=mcp_server) context2.set_session(session2) await context2.progress.report(25, 100) session2.send_progress_notification.assert_not_called() @pytest.mark.asyncio async def test_resource_reading(self, mcp_server): """Test resource reading through context.""" # Mock server's resource reading mcp_server._mcp_read_resource = AsyncMock( return_value=[{"uri": "file://test.txt", "text": "Test content"}] ) context = Context(server=mcp_server) resources = await context.resources.read("file://test.txt") assert len(resources) == 1 assert resources[0]["text"] == "Test content" mcp_server._mcp_read_resource.assert_called_once_with("file://test.txt") @pytest.mark.asyncio async def test_list_roots(self, mcp_server): """Test listing roots.""" session = Mock() # Return an object with roots attribute result = Mock() result.roots = [{"uri": "file:///home", "name": "Home"}] session.list_roots = AsyncMock(return_value=result) context = Context(server=mcp_server) context.set_session(session) roots = await context.resources.list_roots() assert len(roots) == 1 assert roots[0]["name"] == "Home" @pytest.mark.asyncio async def test_sampling(self, mcp_server): """Test sampling functionality.""" session = Mock() # Mock the response with content attribute result = Mock() result.content = {"type": "text", "text": "Response"} session.create_message = AsyncMock(return_value=result) context = Context(server=mcp_server) context.set_session(session) # Mock client capabilities check session.check_client_capability = Mock(return_value=True) # Test basic sampling result = await context.sampling.create_message( messages="Hello", system_prompt="Be helpful", temperature=0.7, max_tokens=100 ) assert result["type"] == "text" assert result["text"] == "Response" # Test with model preferences result = await context.sampling.create_message( messages=[{"role": "user", "content": "Hello"}], model_preferences=ModelPreferences(hints=[ModelHint(name="claude-3")]), ) assert session.create_message.call_count == 2 @pytest.mark.asyncio async def test_sampling_without_capability(self, mcp_server): """Test sampling when client doesn't support it.""" session = Mock() context = Context(server=mcp_server) context.set_session(session) # Mock client capabilities check to return False session.check_client_capability = Mock(return_value=False) with pytest.raises(ValueError) as exc_info: await context.sampling.create_message(messages=["Hello"], max_tokens=32) assert "Client does not support sampling" in str(exc_info.value) @pytest.mark.asyncio async def test_elicitation(self, mcp_server): """Test user input elicitation.""" session = Mock() # Mock the elicit method on session session.elicit = AsyncMock(return_value={"value": "user input"}) context = Context(server=mcp_server) context.set_session(session) # Test string elicitation result = await context.ui.elicit("Enter your name:") assert result == {"value": "user input"} session.elicit.assert_called_once_with( message="Enter your name:", requested_schema={"type": "object", "properties": {}}, timeout=300.0, ) # Test with schema schema = {"type": "object", "properties": {"name": {"type": "string"}}} result = await context.ui.elicit("Enter details:", schema=schema) assert result == {"value": "user input"} assert session.elicit.call_count == 2 session.elicit.assert_called_with( message="Enter details:", requested_schema=schema, timeout=300 ) @pytest.mark.asyncio async def test_notification_queueing(self, mcp_server): """Test notification queueing methods.""" session = Mock() session.send_tool_list_changed = AsyncMock() session.send_resource_list_changed = AsyncMock() session.send_prompt_list_changed = AsyncMock() context = Context(server=mcp_server) context.set_session(session) # Queue notifications - they are queued and not sent immediately await context.notifications.tools.list_changed() await context.notifications.resources.list_changed() await context.notifications.prompts.list_changed() # Notifications should be queued, not sent immediately assert "notifications/tools/list_changed" in context._notification_queue assert "notifications/resources/list_changed" in context._notification_queue assert "notifications/prompts/list_changed" in context._notification_queue # Mock the notification manager nm = Mock() nm.notify_tool_list_changed = AsyncMock() nm.notify_resource_list_changed = AsyncMock() mcp_server.notification_manager = nm # Add session_id to session session.session_id = "test-session-123" # Now flush notifications await context._flush_notifications() # Queue should be cleared after flush assert len(context._notification_queue) == 0 # Verify notifications were sent with the session_id nm.notify_tool_list_changed.assert_called_once_with(["test-session-123"]) nm.notify_resource_list_changed.assert_called_once_with(["test-session-123"]) def test_parse_model_preferences(self, mcp_server): """Test model preferences parsing.""" context = Context(server=mcp_server) # Test with ModelPreferences object prefs = ModelPreferences(hints=[ModelHint(name="gpt-4")]) parsed = context._parse_model_preferences(prefs) assert parsed == prefs # Test with string parsed = context._parse_model_preferences("gpt-4") assert isinstance(parsed, ModelPreferences) assert len(parsed.hints) == 1 assert parsed.hints[0].name == "gpt-4" # Test with list of strings parsed = context._parse_model_preferences(["gpt-4", "claude-3"]) assert isinstance(parsed, ModelPreferences) assert len(parsed.hints) == 2 assert parsed.hints[0].name == "gpt-4" assert parsed.hints[1].name == "claude-3" # Test with None parsed = context._parse_model_preferences(None) assert parsed is None # Test with invalid type with pytest.raises(ValueError, match="Invalid model preferences type"): context._parse_model_preferences({"invalid": "dict"}) @pytest.mark.asyncio async def test_context_without_server(self): """Test operations that require server when server is None.""" # Create a context with a server that will be garbage collected server = Mock() context = Context(server=server) # Clear the strong reference to server del server with pytest.raises(RuntimeError) as exc_info: # This should raise because the weak reference is dead _ = context.server assert "Server instance is no longer available" in str(exc_info.value) @pytest.mark.asyncio async def test_context_without_session(self, mcp_server): """Test operations that require session when session is None.""" context = Context(server=mcp_server) # These should return empty/None without raising roots = await context.resources.list_roots() assert roots == [] # Sampling should raise ValueError when session is None with pytest.raises(ValueError, match="Session not available"): await context.sampling.create_message(messages=["Hello"], max_tokens=32) # Elicit should also raise ValueError when session is None with pytest.raises(ValueError, match="Session not available"): await context.ui.elicit("Enter text") @pytest.mark.asyncio async def test_context_as_context_manager(self, mcp_server): """Test using context as an async context manager.""" context = Context(server=mcp_server) # Enter context async with context as ctx: assert ctx == context # Context should be set as current assert get_current_context() == context # After exit, context should be reset assert get_current_context() is None