Add TaskResultHandler unit tests

Author: maxisbey

tests/experimental/tasks/server/test_task_result_handler.py

@@ -0,0 +1,255 @@
+"""Tests for TaskResultHandler."""
+
+from collections.abc import AsyncIterator
+from typing import Any
+from unittest.mock import AsyncMock, Mock
+
+import anyio
+import pytest
+
+from mcp.server.experimental.task_result_handler import TaskResultHandler
+from mcp.shared.exceptions import McpError
+from mcp.shared.experimental.tasks.in_memory_task_store import InMemoryTaskStore
+from mcp.shared.experimental.tasks.message_queue import InMemoryTaskMessageQueue, QueuedMessage
+from mcp.shared.experimental.tasks.resolver import Resolver
+from mcp.shared.message import SessionMessage
+from mcp.types import (
+    CallToolResult,
+    ErrorData,
+    GetTaskPayloadRequest,
+    GetTaskPayloadRequestParams,
+    GetTaskPayloadResult,
+    JSONRPCRequest,
+    TaskMetadata,
+    TextContent,
+)
+
+
+@pytest.fixture
+async def store() -> AsyncIterator[InMemoryTaskStore]:
+    """Provide a clean store for each test."""
+    s = InMemoryTaskStore()
+    yield s
+    s.cleanup()
+
+
+@pytest.fixture
+def queue() -> InMemoryTaskMessageQueue:
+    """Provide a clean queue for each test."""
+    return InMemoryTaskMessageQueue()
+
+
+@pytest.fixture
+def handler(store: InMemoryTaskStore, queue: InMemoryTaskMessageQueue) -> TaskResultHandler:
+    """Provide a handler for each test."""
+    return TaskResultHandler(store, queue)
+
+
+@pytest.mark.anyio
+async def test_handle_returns_result_for_completed_task(
+    store: InMemoryTaskStore, queue: InMemoryTaskMessageQueue, handler: TaskResultHandler
+) -> None:
+    """Test that handle() returns the stored result for a completed task."""
+    task = await store.create_task(TaskMetadata(ttl=60000), task_id="test-task")
+    result = CallToolResult(content=[TextContent(type="text", text="Done!")])
+    await store.store_result(task.taskId, result)
+    await store.update_task(task.taskId, status="completed")
+
+    mock_session = Mock()
+    mock_session.send_message = AsyncMock()
+
+    request = GetTaskPayloadRequest(params=GetTaskPayloadRequestParams(taskId=task.taskId))
+    response = await handler.handle(request, mock_session, "req-1")
+
+    assert response is not None
+    assert response.meta is not None
+    assert "io.modelcontextprotocol/related-task" in response.meta
+
+
+@pytest.mark.anyio
+async def test_handle_raises_for_nonexistent_task(
+    store: InMemoryTaskStore, queue: InMemoryTaskMessageQueue, handler: TaskResultHandler
+) -> None:
+    """Test that handle() raises McpError for nonexistent task."""
+    mock_session = Mock()
+    request = GetTaskPayloadRequest(params=GetTaskPayloadRequestParams(taskId="nonexistent"))
+
+    with pytest.raises(McpError) as exc_info:
+        await handler.handle(request, mock_session, "req-1")
+
+    assert "not found" in exc_info.value.error.message
+
+
+@pytest.mark.anyio
+async def test_handle_returns_empty_result_when_no_result_stored(
+    store: InMemoryTaskStore, queue: InMemoryTaskMessageQueue, handler: TaskResultHandler
+) -> None:
+    """Test that handle() returns minimal result when task completed without stored result."""
+    task = await store.create_task(TaskMetadata(ttl=60000), task_id="test-task")
+    await store.update_task(task.taskId, status="completed")
+
+    mock_session = Mock()
+    mock_session.send_message = AsyncMock()
+
+    request = GetTaskPayloadRequest(params=GetTaskPayloadRequestParams(taskId=task.taskId))
+    response = await handler.handle(request, mock_session, "req-1")
+
+    assert response is not None
+    assert response.meta is not None
+    assert "io.modelcontextprotocol/related-task" in response.meta
+
+
+@pytest.mark.anyio
+async def test_handle_delivers_queued_messages(
+    store: InMemoryTaskStore, queue: InMemoryTaskMessageQueue, handler: TaskResultHandler
+) -> None:
+    """Test that handle() delivers queued messages before returning."""
+    task = await store.create_task(TaskMetadata(ttl=60000), task_id="test-task")
+
+    queued_msg = QueuedMessage(
+        type="notification",
+        message=JSONRPCRequest(
+            jsonrpc="2.0",
+            id="notif-1",
+            method="test/notification",
+            params={},
+        ),
+    )
+    await queue.enqueue(task.taskId, queued_msg)
+    await store.update_task(task.taskId, status="completed")
+
+    sent_messages: list[SessionMessage] = []
+
+    async def track_send(msg: SessionMessage) -> None:
+        sent_messages.append(msg)
+
+    mock_session = Mock()
+    mock_session.send_message = track_send
+
+    request = GetTaskPayloadRequest(params=GetTaskPayloadRequestParams(taskId=task.taskId))
+    await handler.handle(request, mock_session, "req-1")
+
+    assert len(sent_messages) == 1
+
+
+@pytest.mark.anyio
+async def test_handle_waits_for_task_completion(
+    store: InMemoryTaskStore, queue: InMemoryTaskMessageQueue, handler: TaskResultHandler
+) -> None:
+    """Test that handle() waits for task to complete before returning."""
+    task = await store.create_task(TaskMetadata(ttl=60000), task_id="test-task")
+
+    mock_session = Mock()
+    mock_session.send_message = AsyncMock()
+
+    request = GetTaskPayloadRequest(params=GetTaskPayloadRequestParams(taskId=task.taskId))
+    result_holder: list[GetTaskPayloadResult | None] = [None]
+
+    async def run_handle() -> None:
+        result_holder[0] = await handler.handle(request, mock_session, "req-1")
+
+    async with anyio.create_task_group() as tg:
+        tg.start_soon(run_handle)
+        await anyio.sleep(0.05)
+
+        await store.store_result(task.taskId, CallToolResult(content=[TextContent(type="text", text="Done")]))
+        await store.update_task(task.taskId, status="completed")
+
+    assert result_holder[0] is not None
+
+
+@pytest.mark.anyio
+async def test_route_response_resolves_pending_request(
+    store: InMemoryTaskStore, queue: InMemoryTaskMessageQueue, handler: TaskResultHandler
+) -> None:
+    """Test that route_response() resolves a pending request."""
+    resolver: Resolver[dict[str, Any]] = Resolver()
+    handler._pending_requests["req-123"] = resolver
+
+    result = handler.route_response("req-123", {"status": "ok"})
+
+    assert result is True
+    assert resolver.done()
+    assert await resolver.wait() == {"status": "ok"}
+
+
+@pytest.mark.anyio
+async def test_route_response_returns_false_for_unknown_request(
+    store: InMemoryTaskStore, queue: InMemoryTaskMessageQueue, handler: TaskResultHandler
+) -> None:
+    """Test that route_response() returns False for unknown request ID."""
+    result = handler.route_response("unknown-req", {"status": "ok"})
+    assert result is False
+
+
+@pytest.mark.anyio
+async def test_route_response_returns_false_for_already_done_resolver(
+    store: InMemoryTaskStore, queue: InMemoryTaskMessageQueue, handler: TaskResultHandler
+) -> None:
+    """Test that route_response() returns False if resolver already completed."""
+    resolver: Resolver[dict[str, Any]] = Resolver()
+    resolver.set_result({"already": "done"})
+    handler._pending_requests["req-123"] = resolver
+
+    result = handler.route_response("req-123", {"new": "data"})
+
+    assert result is False
+
+
+@pytest.mark.anyio
+async def test_route_error_resolves_pending_request_with_exception(
+    store: InMemoryTaskStore, queue: InMemoryTaskMessageQueue, handler: TaskResultHandler
+) -> None:
+    """Test that route_error() sets exception on pending request."""
+    resolver: Resolver[dict[str, Any]] = Resolver()
+    handler._pending_requests["req-123"] = resolver
+
+    error = ErrorData(code=-32600, message="Something went wrong")
+    result = handler.route_error("req-123", error)
+
+    assert result is True
+    assert resolver.done()
+
+    with pytest.raises(McpError) as exc_info:
+        await resolver.wait()
+    assert exc_info.value.error.message == "Something went wrong"
+
+
+@pytest.mark.anyio
+async def test_route_error_returns_false_for_unknown_request(
+    store: InMemoryTaskStore, queue: InMemoryTaskMessageQueue, handler: TaskResultHandler
+) -> None:
+    """Test that route_error() returns False for unknown request ID."""
+    error = ErrorData(code=-32600, message="Error")
+    result = handler.route_error("unknown-req", error)
+    assert result is False
+
+
+@pytest.mark.anyio
+async def test_deliver_registers_resolver_for_request_messages(
+    store: InMemoryTaskStore, queue: InMemoryTaskMessageQueue, handler: TaskResultHandler
+) -> None:
+    """Test that _deliver_queued_messages registers resolvers for request messages."""
+    task = await store.create_task(TaskMetadata(ttl=60000), task_id="test-task")
+
+    resolver: Resolver[dict[str, Any]] = Resolver()
+    queued_msg = QueuedMessage(
+        type="request",
+        message=JSONRPCRequest(
+            jsonrpc="2.0",
+            id="inner-req-1",
+            method="elicitation/create",
+            params={},
+        ),
+        resolver=resolver,
+        original_request_id="inner-req-1",
+    )
+    await queue.enqueue(task.taskId, queued_msg)
+
+    mock_session = Mock()
+    mock_session.send_message = AsyncMock()
+
+    await handler._deliver_queued_messages(task.taskId, mock_session, "outer-req-1")
+
+    assert "inner-req-1" in handler._pending_requests
+    assert handler._pending_requests["inner-req-1"] is resolver