refactor: replace lowlevel Server decorators with on_* constructor kwargs (#1985)

tests/shared/test_ws.py

@@ -2,7 +2,6 @@ import multiprocessing
 import socket
 import time
 from collections.abc import AsyncGenerator, Generator
-from typing import Any
 from urllib.parse import urlparse
 
 import anyio
@@ -15,9 +14,21 @@ from starlette.websockets import WebSocket
 from mcp import MCPError
 from mcp.client.session import ClientSession
 from mcp.client.websocket import websocket_client
-from mcp.server import Server
+from mcp.server import Server, ServerRequestContext
 from mcp.server.websocket import websocket_server
-from mcp.types import EmptyResult, InitializeResult, ReadResourceResult, TextContent, TextResourceContents, Tool
+from mcp.types import (
+    CallToolRequestParams,
+    CallToolResult,
+    EmptyResult,
+    InitializeResult,
+    ListToolsResult,
+    PaginatedRequestParams,
+    ReadResourceRequestParams,
+    ReadResourceResult,
+    TextContent,
+    TextResourceContents,
+    Tool,
+)
 from tests.test_helpers import wait_for_server
 
 SERVER_NAME = "test_server_for_WS"
@@ -35,42 +46,59 @@ def server_url(server_port: int) -> str:
     return f"ws://127.0.0.1:{server_port}"
 
 
-# Test server implementation
-class ServerTest(Server):  # pragma: no cover
-    def __init__(self):
-        super().__init__(SERVER_NAME)
-
-        @self.read_resource()
-        async def handle_read_resource(uri: str) -> str | bytes:
-            parsed = urlparse(uri)
-            if parsed.scheme == "foobar":
-                return f"Read {parsed.netloc}"
-            elif parsed.scheme == "slow":
-                # Simulate a slow resource
-                await anyio.sleep(2.0)
-                return f"Slow response from {parsed.netloc}"
-
-            raise MCPError(code=404, message="OOPS! no resource with that URI was found")
-
-        @self.list_tools()
-        async def handle_list_tools() -> list[Tool]:
-            return [
-                Tool(
-                    name="test_tool",
-                    description="A test tool",
-                    input_schema={"type": "object", "properties": {}},
+async def handle_read_resource(  # pragma: no cover
+    ctx: ServerRequestContext, params: ReadResourceRequestParams
+) -> ReadResourceResult:
+    parsed = urlparse(str(params.uri))
+    if parsed.scheme == "foobar":
+        return ReadResourceResult(
+            contents=[TextResourceContents(uri=str(params.uri), text=f"Read {parsed.netloc}", mime_type="text/plain")]
+        )
+    elif parsed.scheme == "slow":
+        await anyio.sleep(2.0)
+        return ReadResourceResult(
+            contents=[
+                TextResourceContents(
+                    uri=str(params.uri), text=f"Slow response from {parsed.netloc}", mime_type="text/plain"
                 )
             ]
+        )
+    raise MCPError(code=404, message="OOPS! no resource with that URI was found")
 
-        @self.call_tool()
-        async def handle_call_tool(name: str, args: dict[str, Any]) -> list[TextContent]:
-            return [TextContent(type="text", text=f"Called {name}")]
+
+async def handle_list_tools(  # pragma: no cover
+    ctx: ServerRequestContext, params: PaginatedRequestParams | None
+) -> ListToolsResult:
+    return ListToolsResult(
+        tools=[
+            Tool(
+                name="test_tool",
+                description="A test tool",
+                input_schema={"type": "object", "properties": {}},
+            )
+        ]
+    )
+
+
+async def handle_call_tool(  # pragma: no cover
+    ctx: ServerRequestContext, params: CallToolRequestParams
+) -> CallToolResult:
+    return CallToolResult(content=[TextContent(type="text", text=f"Called {params.name}")])
+
+
+def _create_server() -> Server:  # pragma: no cover
+    return Server(
+        SERVER_NAME,
+        on_read_resource=handle_read_resource,
+        on_list_tools=handle_list_tools,
+        on_call_tool=handle_call_tool,
+    )
 
 
 # Test fixtures
 def make_server_app() -> Starlette:  # pragma: no cover
     """Create test Starlette app with WebSocket transport"""
-    server = ServerTest()
+    server = _create_server()
 
     async def handle_ws(websocket: WebSocket):
         async with websocket_server(websocket.scope, websocket.receive, websocket.send) as streams: