Add comprehensive capability tests and improve test coverage

maxisbey · maxisbey · commit b7d44fae72fe · 2025-11-28T15:23:10.000Z
- Add test_capabilities.py with unit tests for all capability checking functions
- Add tests for elicit_as_task and create_message_as_task without handler
- Add scenario 4 sampling test (task-augmented tool call + task-augmented sampling)
- Replace sleep-based polling with event-based synchronization for faster, deterministic tests
- Simplify for/else patterns in test code
- Add additional check_tasks_capability edge case tests

Test coverage improved to 99.94% with 0 missing statements.
diff --git a/tests/experimental/tasks/server/test_server_task_context.py b/tests/experimental/tasks/server/test_server_task_context.py
@@ -502,7 +502,7 @@ async def test_create_message_restores_status_on_cancellation() -> None:
     """Test that create_message() restores task status to working when cancelled."""
     import anyio
 
-    from mcp.types import JSONRPCRequest, SamplingMessage, TextContent
+    from mcp.types import JSONRPCRequest, SamplingMessage
 
     store = InMemoryTaskStore()
     queue = InMemoryTaskMessageQueue()
@@ -570,3 +570,98 @@ async def do_sampling() -> None:
     assert cancelled_error_raised
 
     store.cleanup()
+
+
+@pytest.mark.anyio
+async def test_elicit_as_task_raises_without_handler() -> None:
+    """Test that elicit_as_task() raises when handler is not provided."""
+    from mcp.types import (
+        ClientCapabilities,
+        ClientTasksCapability,
+        ClientTasksRequestsCapability,
+        Implementation,
+        InitializeRequestParams,
+        TasksCreateElicitationCapability,
+        TasksElicitationCapability,
+    )
+
+    store = InMemoryTaskStore()
+    queue = InMemoryTaskMessageQueue()
+    task = await store.create_task(TaskMetadata(ttl=60000))
+
+    # Create mock session with proper client capabilities
+    mock_session = Mock()
+    mock_session.client_params = InitializeRequestParams(
+        protocolVersion="2025-01-01",
+        capabilities=ClientCapabilities(
+            tasks=ClientTasksCapability(
+                requests=ClientTasksRequestsCapability(
+                    elicitation=TasksElicitationCapability(create=TasksCreateElicitationCapability())
+                )
+            )
+        ),
+        clientInfo=Implementation(name="test", version="1.0"),
+    )
+
+    ctx = ServerTaskContext(
+        task=task,
+        store=store,
+        session=mock_session,
+        queue=queue,
+        handler=None,  # No handler
+    )
+
+    with pytest.raises(RuntimeError, match="handler is required for elicit_as_task"):
+        await ctx.elicit_as_task(message="Test?", requestedSchema={"type": "object"})
+
+    store.cleanup()
+
+
+@pytest.mark.anyio
+async def test_create_message_as_task_raises_without_handler() -> None:
+    """Test that create_message_as_task() raises when handler is not provided."""
+    from mcp.types import (
+        ClientCapabilities,
+        ClientTasksCapability,
+        ClientTasksRequestsCapability,
+        Implementation,
+        InitializeRequestParams,
+        SamplingMessage,
+        TasksCreateMessageCapability,
+        TasksSamplingCapability,
+        TextContent,
+    )
+
+    store = InMemoryTaskStore()
+    queue = InMemoryTaskMessageQueue()
+    task = await store.create_task(TaskMetadata(ttl=60000))
+
+    # Create mock session with proper client capabilities
+    mock_session = Mock()
+    mock_session.client_params = InitializeRequestParams(
+        protocolVersion="2025-01-01",
+        capabilities=ClientCapabilities(
+            tasks=ClientTasksCapability(
+                requests=ClientTasksRequestsCapability(
+                    sampling=TasksSamplingCapability(createMessage=TasksCreateMessageCapability())
+                )
+            )
+        ),
+        clientInfo=Implementation(name="test", version="1.0"),
+    )
+
+    ctx = ServerTaskContext(
+        task=task,
+        store=store,
+        session=mock_session,
+        queue=queue,
+        handler=None,  # No handler
+    )
+
+    with pytest.raises(RuntimeError, match="handler is required for create_message_as_task"):
+        await ctx.create_message_as_task(
+            messages=[SamplingMessage(role="user", content=TextContent(type="text", text="Hello"))],
+            max_tokens=100,
+        )
+
+    store.cleanup()
diff --git a/tests/experimental/tasks/test_capabilities.py b/tests/experimental/tasks/test_capabilities.py
@@ -0,0 +1,283 @@
+"""Tests for tasks capability checking utilities."""
+
+import pytest
+
+from mcp.shared.exceptions import McpError
+from mcp.shared.experimental.tasks.capabilities import (
+    check_tasks_capability,
+    has_task_augmented_elicitation,
+    has_task_augmented_sampling,
+    require_task_augmented_elicitation,
+    require_task_augmented_sampling,
+)
+from mcp.types import (
+    ClientCapabilities,
+    ClientTasksCapability,
+    ClientTasksRequestsCapability,
+    TasksCreateElicitationCapability,
+    TasksCreateMessageCapability,
+    TasksElicitationCapability,
+    TasksSamplingCapability,
+)
+
+
+class TestCheckTasksCapability:
+    """Tests for check_tasks_capability function."""
+
+    def test_required_requests_none_returns_true(self) -> None:
+        """When required.requests is None, should return True."""
+        required = ClientTasksCapability()
+        client = ClientTasksCapability()
+        assert check_tasks_capability(required, client) is True
+
+    def test_client_requests_none_returns_false(self) -> None:
+        """When client.requests is None but required.requests is set, should return False."""
+        required = ClientTasksCapability(requests=ClientTasksRequestsCapability())
+        client = ClientTasksCapability()
+        assert check_tasks_capability(required, client) is False
+
+    def test_elicitation_required_but_client_missing(self) -> None:
+        """When elicitation is required but client doesn't have it."""
+        required = ClientTasksCapability(
+            requests=ClientTasksRequestsCapability(elicitation=TasksElicitationCapability())
+        )
+        client = ClientTasksCapability(requests=ClientTasksRequestsCapability())
+        assert check_tasks_capability(required, client) is False
+
+    def test_elicitation_create_required_but_client_missing(self) -> None:
+        """When elicitation.create is required but client doesn't have it."""
+        required = ClientTasksCapability(
+            requests=ClientTasksRequestsCapability(
+                elicitation=TasksElicitationCapability(create=TasksCreateElicitationCapability())
+            )
+        )
+        client = ClientTasksCapability(
+            requests=ClientTasksRequestsCapability(
+                elicitation=TasksElicitationCapability()  # No create
+            )
+        )
+        assert check_tasks_capability(required, client) is False
+
+    def test_elicitation_create_present(self) -> None:
+        """When elicitation.create is required and client has it."""
+        required = ClientTasksCapability(
+            requests=ClientTasksRequestsCapability(
+                elicitation=TasksElicitationCapability(create=TasksCreateElicitationCapability())
+            )
+        )
+        client = ClientTasksCapability(
+            requests=ClientTasksRequestsCapability(
+                elicitation=TasksElicitationCapability(create=TasksCreateElicitationCapability())
+            )
+        )
+        assert check_tasks_capability(required, client) is True
+
+    def test_sampling_required_but_client_missing(self) -> None:
+        """When sampling is required but client doesn't have it."""
+        required = ClientTasksCapability(requests=ClientTasksRequestsCapability(sampling=TasksSamplingCapability()))
+        client = ClientTasksCapability(requests=ClientTasksRequestsCapability())
+        assert check_tasks_capability(required, client) is False
+
+    def test_sampling_create_message_required_but_client_missing(self) -> None:
+        """When sampling.createMessage is required but client doesn't have it."""
+        required = ClientTasksCapability(
+            requests=ClientTasksRequestsCapability(
+                sampling=TasksSamplingCapability(createMessage=TasksCreateMessageCapability())
+            )
+        )
+        client = ClientTasksCapability(
+            requests=ClientTasksRequestsCapability(
+                sampling=TasksSamplingCapability()  # No createMessage
+            )
+        )
+        assert check_tasks_capability(required, client) is False
+
+    def test_sampling_create_message_present(self) -> None:
+        """When sampling.createMessage is required and client has it."""
+        required = ClientTasksCapability(
+            requests=ClientTasksRequestsCapability(
+                sampling=TasksSamplingCapability(createMessage=TasksCreateMessageCapability())
+            )
+        )
+        client = ClientTasksCapability(
+            requests=ClientTasksRequestsCapability(
+                sampling=TasksSamplingCapability(createMessage=TasksCreateMessageCapability())
+            )
+        )
+        assert check_tasks_capability(required, client) is True
+
+    def test_both_elicitation_and_sampling_present(self) -> None:
+        """When both elicitation.create and sampling.createMessage are required and client has both."""
+        required = ClientTasksCapability(
+            requests=ClientTasksRequestsCapability(
+                elicitation=TasksElicitationCapability(create=TasksCreateElicitationCapability()),
+                sampling=TasksSamplingCapability(createMessage=TasksCreateMessageCapability()),
+            )
+        )
+        client = ClientTasksCapability(
+            requests=ClientTasksRequestsCapability(
+                elicitation=TasksElicitationCapability(create=TasksCreateElicitationCapability()),
+                sampling=TasksSamplingCapability(createMessage=TasksCreateMessageCapability()),
+            )
+        )
+        assert check_tasks_capability(required, client) is True
+
+    def test_elicitation_without_create_required(self) -> None:
+        """When elicitation is required but not create specifically."""
+        required = ClientTasksCapability(
+            requests=ClientTasksRequestsCapability(
+                elicitation=TasksElicitationCapability()  # No create
+            )
+        )
+        client = ClientTasksCapability(
+            requests=ClientTasksRequestsCapability(
+                elicitation=TasksElicitationCapability(create=TasksCreateElicitationCapability())
+            )
+        )
+        assert check_tasks_capability(required, client) is True
+
+    def test_sampling_without_create_message_required(self) -> None:
+        """When sampling is required but not createMessage specifically."""
+        required = ClientTasksCapability(
+            requests=ClientTasksRequestsCapability(
+                sampling=TasksSamplingCapability()  # No createMessage
+            )
+        )
+        client = ClientTasksCapability(
+            requests=ClientTasksRequestsCapability(
+                sampling=TasksSamplingCapability(createMessage=TasksCreateMessageCapability())
+            )
+        )
+        assert check_tasks_capability(required, client) is True
+
+
+class TestHasTaskAugmentedElicitation:
+    """Tests for has_task_augmented_elicitation function."""
+
+    def test_tasks_none(self) -> None:
+        """Returns False when caps.tasks is None."""
+        caps = ClientCapabilities()
+        assert has_task_augmented_elicitation(caps) is False
+
+    def test_requests_none(self) -> None:
+        """Returns False when caps.tasks.requests is None."""
+        caps = ClientCapabilities(tasks=ClientTasksCapability())
+        assert has_task_augmented_elicitation(caps) is False
+
+    def test_elicitation_none(self) -> None:
+        """Returns False when caps.tasks.requests.elicitation is None."""
+        caps = ClientCapabilities(tasks=ClientTasksCapability(requests=ClientTasksRequestsCapability()))
+        assert has_task_augmented_elicitation(caps) is False
+
+    def test_create_none(self) -> None:
+        """Returns False when caps.tasks.requests.elicitation.create is None."""
+        caps = ClientCapabilities(
+            tasks=ClientTasksCapability(
+                requests=ClientTasksRequestsCapability(elicitation=TasksElicitationCapability())
+            )
+        )
+        assert has_task_augmented_elicitation(caps) is False
+
+    def test_create_present(self) -> None:
+        """Returns True when full capability path is present."""
+        caps = ClientCapabilities(
+            tasks=ClientTasksCapability(
+                requests=ClientTasksRequestsCapability(
+                    elicitation=TasksElicitationCapability(create=TasksCreateElicitationCapability())
+                )
+            )
+        )
+        assert has_task_augmented_elicitation(caps) is True
+
+
+class TestHasTaskAugmentedSampling:
+    """Tests for has_task_augmented_sampling function."""
+
+    def test_tasks_none(self) -> None:
+        """Returns False when caps.tasks is None."""
+        caps = ClientCapabilities()
+        assert has_task_augmented_sampling(caps) is False
+
+    def test_requests_none(self) -> None:
+        """Returns False when caps.tasks.requests is None."""
+        caps = ClientCapabilities(tasks=ClientTasksCapability())
+        assert has_task_augmented_sampling(caps) is False
+
+    def test_sampling_none(self) -> None:
+        """Returns False when caps.tasks.requests.sampling is None."""
+        caps = ClientCapabilities(tasks=ClientTasksCapability(requests=ClientTasksRequestsCapability()))
+        assert has_task_augmented_sampling(caps) is False
+
+    def test_create_message_none(self) -> None:
+        """Returns False when caps.tasks.requests.sampling.createMessage is None."""
+        caps = ClientCapabilities(
+            tasks=ClientTasksCapability(requests=ClientTasksRequestsCapability(sampling=TasksSamplingCapability()))
+        )
+        assert has_task_augmented_sampling(caps) is False
+
+    def test_create_message_present(self) -> None:
+        """Returns True when full capability path is present."""
+        caps = ClientCapabilities(
+            tasks=ClientTasksCapability(
+                requests=ClientTasksRequestsCapability(
+                    sampling=TasksSamplingCapability(createMessage=TasksCreateMessageCapability())
+                )
+            )
+        )
+        assert has_task_augmented_sampling(caps) is True
+
+
+class TestRequireTaskAugmentedElicitation:
+    """Tests for require_task_augmented_elicitation function."""
+
+    def test_raises_when_none(self) -> None:
+        """Raises McpError when client_caps is None."""
+        with pytest.raises(McpError) as exc_info:
+            require_task_augmented_elicitation(None)
+        assert "task-augmented elicitation" in str(exc_info.value)
+
+    def test_raises_when_missing(self) -> None:
+        """Raises McpError when capability is missing."""
+        caps = ClientCapabilities()
+        with pytest.raises(McpError) as exc_info:
+            require_task_augmented_elicitation(caps)
+        assert "task-augmented elicitation" in str(exc_info.value)
+
+    def test_passes_when_present(self) -> None:
+        """Does not raise when capability is present."""
+        caps = ClientCapabilities(
+            tasks=ClientTasksCapability(
+                requests=ClientTasksRequestsCapability(
+                    elicitation=TasksElicitationCapability(create=TasksCreateElicitationCapability())
+                )
+            )
+        )
+        require_task_augmented_elicitation(caps)  # Should not raise
+
+
+class TestRequireTaskAugmentedSampling:
+    """Tests for require_task_augmented_sampling function."""
+
+    def test_raises_when_none(self) -> None:
+        """Raises McpError when client_caps is None."""
+        with pytest.raises(McpError) as exc_info:
+            require_task_augmented_sampling(None)
+        assert "task-augmented sampling" in str(exc_info.value)
+
+    def test_raises_when_missing(self) -> None:
+        """Raises McpError when capability is missing."""
+        caps = ClientCapabilities()
+        with pytest.raises(McpError) as exc_info:
+            require_task_augmented_sampling(caps)
+        assert "task-augmented sampling" in str(exc_info.value)
+
+    def test_passes_when_present(self) -> None:
+        """Does not raise when capability is present."""
+        caps = ClientCapabilities(
+            tasks=ClientTasksCapability(
+                requests=ClientTasksRequestsCapability(
+                    sampling=TasksSamplingCapability(createMessage=TasksCreateMessageCapability())
+                )
+            )
+        )
+        require_task_augmented_sampling(caps)  # Should not raise
diff --git a/tests/experimental/tasks/test_elicitation_scenarios.py b/tests/experimental/tasks/test_elicitation_scenarios.py
