Clean up test code patterns

- Remove section header comments from test files
- Move all inline imports to top of files
- Replace hardcoded error codes (-32600) with INVALID_REQUEST constant
- Replace arbitrary sleeps with polling loops for deterministic tests
- Add pragma no branch to polling conditions that always succeed on first try

Author: maxisbey

tests/experimental/tasks/server/test_context.py

@@ -7,8 +7,6 @@
 from mcp.shared.experimental.tasks.in_memory_task_store import InMemoryTaskStore
 from mcp.types import CallToolResult, TaskMetadata, TextContent
 
-# --- TaskContext tests ---
-
 
 @pytest.mark.anyio
 async def test_task_context_properties() -> None:
@@ -98,9 +96,6 @@ async def test_task_context_cancellation() -> None:
     store.cleanup()
 
 
-# --- create_task_state tests ---
-
-
 def test_create_task_state_generates_id() -> None:
     """create_task_state generates a unique task ID when none provided."""
     task1 = create_task_state(TaskMetadata(ttl=60000))
@@ -127,9 +122,6 @@ def test_create_task_state_has_created_at() -> None:
     assert task.createdAt is not None
 
 
-# --- task_execution context manager tests ---
-
-
 @pytest.mark.anyio
 async def test_task_execution_provides_context() -> None:
     """task_execution provides a TaskContext for the task."""

tests/experimental/tasks/server/test_run_task_flow.py

@@ -10,20 +10,33 @@
 """
 
 from typing import Any
+from unittest.mock import Mock
 
 import anyio
 import pytest
 from anyio import Event
 
 from mcp.client.session import ClientSession
 from mcp.server import Server
+from mcp.server.experimental.request_context import Experimental
 from mcp.server.experimental.task_context import ServerTaskContext
+from mcp.server.experimental.task_support import TaskSupport
 from mcp.server.lowlevel import NotificationOptions
+from mcp.shared.experimental.tasks.in_memory_task_store import InMemoryTaskStore
+from mcp.shared.experimental.tasks.message_queue import InMemoryTaskMessageQueue
 from mcp.shared.message import SessionMessage
 from mcp.types import (
     TASK_REQUIRED,
     CallToolResult,
+    CancelTaskRequest,
+    CancelTaskResult,
     CreateTaskResult,
+    GetTaskPayloadRequest,
+    GetTaskPayloadResult,
+    GetTaskRequest,
+    GetTaskResult,
+    ListTasksRequest,
+    ListTasksResult,
     TextContent,
     Tool,
     ToolExecution,
@@ -113,12 +126,12 @@ async def run_client() -> None:
             with anyio.fail_after(5):
                 await work_completed.wait()
 
-            # Small delay to let task state update
-            await anyio.sleep(0.1)
-
-            # Poll task status
-            task_status = await client_session.experimental.get_task(task_id)
-            assert task_status.status == "completed"
+            # Poll until task status is completed
+            with anyio.fail_after(5):
+                while True:
+                    task_status = await client_session.experimental.get_task(task_id)
+                    if task_status.status == "completed":  # pragma: no branch
+                        break
 
     async with anyio.create_task_group() as tg:
         tg.start_soon(run_server)
@@ -181,11 +194,13 @@ async def run_client() -> None:
             with anyio.fail_after(5):
                 await work_failed.wait()
 
-            await anyio.sleep(0.1)
+            # Poll until task status is failed
+            with anyio.fail_after(5):
+                while True:
+                    task_status = await client_session.experimental.get_task(task_id)
+                    if task_status.status == "failed":  # pragma: no branch
+                        break
 
-            # Task should be failed
-            task_status = await client_session.experimental.get_task(task_id)
-            assert task_status.status == "failed"
             assert "Something went wrong" in (task_status.statusMessage or "")
 
     async with anyio.create_task_group() as tg:
@@ -217,9 +232,6 @@ async def test_enable_tasks_auto_registers_handlers() -> None:
 @pytest.mark.anyio
 async def test_enable_tasks_with_custom_store_and_queue() -> None:
     """Test that enable_tasks() uses provided store and queue instead of defaults."""
-    from mcp.shared.experimental.tasks.in_memory_task_store import InMemoryTaskStore
-    from mcp.shared.experimental.tasks.message_queue import InMemoryTaskMessageQueue
-
     server = Server("test-custom-store-queue")
 
     # Create custom store and queue
@@ -237,17 +249,6 @@ async def test_enable_tasks_with_custom_store_and_queue() -> None:
 @pytest.mark.anyio
 async def test_enable_tasks_skips_default_handlers_when_custom_registered() -> None:
     """Test that enable_tasks() doesn't override already-registered handlers."""
-    from mcp.types import (
-        CancelTaskRequest,
-        CancelTaskResult,
-        GetTaskPayloadRequest,
-        GetTaskPayloadResult,
-        GetTaskRequest,
-        GetTaskResult,
-        ListTasksRequest,
-        ListTasksResult,
-    )
-
     server = Server("test-custom-handlers")
 
     # Register custom handlers BEFORE enable_tasks (never called, just for registration)
@@ -281,8 +282,6 @@ async def custom_cancel_task(req: CancelTaskRequest) -> CancelTaskResult:
 @pytest.mark.anyio
 async def test_run_task_without_enable_tasks_raises() -> None:
     """Test that run_task raises when enable_tasks() wasn't called."""
-    from mcp.server.experimental.request_context import Experimental
-
     experimental = Experimental(
         task_metadata=None,
         _client_capabilities=None,
@@ -300,8 +299,6 @@ async def work(task: ServerTaskContext) -> CallToolResult:
 @pytest.mark.anyio
 async def test_task_support_task_group_before_run_raises() -> None:
     """Test that accessing task_group before run() raises RuntimeError."""
-    from mcp.server.experimental.task_support import TaskSupport
-
     task_support = TaskSupport.in_memory()
 
     with pytest.raises(RuntimeError, match="TaskSupport not running"):
@@ -311,9 +308,6 @@ async def test_task_support_task_group_before_run_raises() -> None:
 @pytest.mark.anyio
 async def test_run_task_without_session_raises() -> None:
     """Test that run_task raises when session is not available."""
-    from mcp.server.experimental.request_context import Experimental
-    from mcp.server.experimental.task_support import TaskSupport
-
     task_support = TaskSupport.in_memory()
 
     experimental = Experimental(
@@ -333,11 +327,6 @@ async def work(task: ServerTaskContext) -> CallToolResult:
 @pytest.mark.anyio
 async def test_run_task_without_task_metadata_raises() -> None:
     """Test that run_task raises when request is not task-augmented."""
-    from unittest.mock import Mock
-
-    from mcp.server.experimental.request_context import Experimental
-    from mcp.server.experimental.task_support import TaskSupport
-
     task_support = TaskSupport.in_memory()
     mock_session = Mock()
 
@@ -469,11 +458,12 @@ async def run_client() -> None:
             with anyio.fail_after(5):
                 await work_completed.wait()
 
-            await anyio.sleep(0.1)
-
-            # Task should be completed (from manual complete, not auto-complete)
-            status = await client_session.experimental.get_task(task_id)
-            assert status.status == "completed"
+            # Poll until task status is completed
+            with anyio.fail_after(5):
+                while True:
+                    status = await client_session.experimental.get_task(task_id)
+                    if status.status == "completed":  # pragma: no branch
+                        break
 
     async with anyio.create_task_group() as tg:
         tg.start_soon(run_server)
@@ -533,11 +523,14 @@ async def run_client() -> None:
             with anyio.fail_after(5):
                 await work_completed.wait()
 
-            await anyio.sleep(0.1)
+            # Poll until task status is failed
+            with anyio.fail_after(5):
+                while True:
+                    status = await client_session.experimental.get_task(task_id)
+                    if status.status == "failed":  # pragma: no branch
+                        break
 
             # Task should still be failed (from manual fail, not auto-fail from exception)
-            status = await client_session.experimental.get_task(task_id)
-            assert status.status == "failed"
             assert status.statusMessage == "Manually failed"  # Not "This error should not change status"
 
     async with anyio.create_task_group() as tg:

tests/experimental/tasks/server/test_server.py

@@ -8,12 +8,18 @@
 
 from mcp.client.session import ClientSession
 from mcp.server import Server
+from mcp.server.experimental.task_result_handler import TaskResultHandler
 from mcp.server.lowlevel import NotificationOptions
 from mcp.server.models import InitializationOptions
 from mcp.server.session import ServerSession
-from mcp.shared.message import SessionMessage
+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
+from mcp.shared.message import ServerMessageMetadata, SessionMessage
+from mcp.shared.response_router import ResponseRouter
 from mcp.shared.session import RequestResponder
 from mcp.types import (
+    INVALID_REQUEST,
     TASK_FORBIDDEN,
     TASK_OPTIONAL,
     TASK_REQUIRED,
@@ -52,8 +58,6 @@
     ToolExecution,
 )
 
-# --- Experimental handler tests ---
-
 
 @pytest.mark.anyio
 async def test_list_tasks_handler() -> None:
@@ -174,9 +178,6 @@ async def handle_cancel_task(request: CancelTaskRequest) -> CancelTaskResult:
     assert result.root.status == "cancelled"
 
 
-# --- Server capabilities tests ---
-
-
 @pytest.mark.anyio
 async def test_server_capabilities_include_tasks() -> None:
     """Test that server capabilities include tasks when handlers are registered."""
@@ -223,9 +224,6 @@ async def handle_list_tasks(request: ListTasksRequest) -> ListTasksResult:
     assert capabilities.tasks.cancel is None  # Not registered
 
 
-# --- Tool annotation tests ---
-
-
 @pytest.mark.anyio
 async def test_tool_with_task_execution_metadata() -> None:
     """Test that tools can declare task execution mode."""
@@ -270,9 +268,6 @@ async def list_tools():
     assert tools[2].execution.taskSupport == TASK_OPTIONAL
 
 
-# --- Integration tests ---
-
-
 @pytest.mark.anyio
 async def test_task_metadata_in_call_tool_request() -> None:
     """Test that task metadata is accessible via RequestContext when calling a tool."""
@@ -471,8 +466,6 @@ async def test_default_task_handlers_via_enable_tasks() -> None:
     - _default_list_tasks
     - _default_cancel_task
     """
-    from mcp.shared.exceptions import McpError
-
     server = Server("test-default-handlers")
     # Enable tasks with default handlers (no custom handlers registered)
     task_support = server.experimental.enable_tasks()
@@ -569,10 +562,6 @@ async def run_server() -> None:
 @pytest.mark.anyio
 async def test_set_task_result_handler() -> None:
     """Test that set_task_result_handler adds the handler as a response router."""
-    from mcp.server.experimental.task_result_handler import TaskResultHandler
-    from mcp.shared.experimental.tasks.in_memory_task_store import InMemoryTaskStore
-    from mcp.shared.experimental.tasks.message_queue import InMemoryTaskMessageQueue
-
     server_to_client_send, server_to_client_receive = anyio.create_memory_object_stream[SessionMessage](10)
     client_to_server_send, client_to_server_receive = anyio.create_memory_object_stream[SessionMessage](10)
 
@@ -751,8 +740,6 @@ async def test_build_create_message_request() -> None:
 @pytest.mark.anyio
 async def test_send_message() -> None:
     """Test that send_message sends a raw session message."""
-    from mcp.shared.message import ServerMessageMetadata
-
     server_to_client_send, server_to_client_receive = anyio.create_memory_object_stream[SessionMessage](10)
     client_to_server_send, client_to_server_receive = anyio.create_memory_object_stream[SessionMessage](10)
 
@@ -790,8 +777,6 @@ async def test_send_message() -> None:
 @pytest.mark.anyio
 async def test_response_routing_success() -> None:
     """Test that response routing works for success responses."""
-    from mcp.shared.session import ResponseRouter
-
     server_to_client_send, server_to_client_receive = anyio.create_memory_object_stream[SessionMessage](10)
     client_to_server_send, client_to_server_receive = anyio.create_memory_object_stream[SessionMessage](10)
 
@@ -846,8 +831,6 @@ def route_error(self, request_id: str | int, error: ErrorData) -> bool:
 @pytest.mark.anyio
 async def test_response_routing_error() -> None:
     """Test that error routing works for error responses."""
-    from mcp.shared.session import ResponseRouter
-
     server_to_client_send, server_to_client_receive = anyio.create_memory_object_stream[SessionMessage](10)
     client_to_server_send, client_to_server_receive = anyio.create_memory_object_stream[SessionMessage](10)
 
@@ -878,7 +861,7 @@ def route_error(self, request_id: str | int, error: ErrorData) -> bool:
             server_session.add_response_router(router)
 
             # Simulate receiving an error response from client
-            error_data = ErrorData(code=-32600, message="Test error")
+            error_data = ErrorData(code=INVALID_REQUEST, message="Test error")
             error_response = JSONRPCError(jsonrpc="2.0", id="test-req-2", error=error_data)
             message = SessionMessage(message=JSONRPCMessage(error_response))
 
@@ -903,8 +886,6 @@ def route_error(self, request_id: str | int, error: ErrorData) -> bool:
 @pytest.mark.anyio
 async def test_response_routing_skips_non_matching_routers() -> None:
     """Test that routing continues to next router when first doesn't match."""
-    from mcp.shared.session import ResponseRouter
-
     server_to_client_send, server_to_client_receive = anyio.create_memory_object_stream[SessionMessage](10)
     client_to_server_send, client_to_server_receive = anyio.create_memory_object_stream[SessionMessage](10)
 
@@ -963,8 +944,6 @@ def route_error(self, request_id: str | int, error: ErrorData) -> bool:
 @pytest.mark.anyio
 async def test_error_routing_skips_non_matching_routers() -> None:
     """Test that error routing continues to next router when first doesn't match."""
-    from mcp.shared.session import ResponseRouter
-
     server_to_client_send, server_to_client_receive = anyio.create_memory_object_stream[SessionMessage](10)
     client_to_server_send, client_to_server_receive = anyio.create_memory_object_stream[SessionMessage](10)
 
@@ -1004,7 +983,7 @@ def route_error(self, request_id: str | int, error: ErrorData) -> bool:
             server_session.add_response_router(MatchingRouter())
 
             # Send an error - should skip first router and be handled by second
-            error_data = ErrorData(code=-32600, message="Test error")
+            error_data = ErrorData(code=INVALID_REQUEST, message="Test error")
             error_response = JSONRPCError(jsonrpc="2.0", id="test-req-2", error=error_data)
             message = SessionMessage(message=JSONRPCMessage(error_response))
             await client_to_server_send.send(message)