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

tests/server/test_completion_with_context.py

@@ -1,15 +1,13 @@
 """Tests for completion handler with context functionality."""
 
-from typing import Any
-
 import pytest
 
 from mcp import Client
-from mcp.server.lowlevel import Server
+from mcp.server import Server, ServerRequestContext
 from mcp.types import (
+    CompleteRequestParams,
+    CompleteResult,
     Completion,
-    CompletionArgument,
-    CompletionContext,
     PromptReference,
     ResourceTemplateReference,
 )
@@ -18,23 +16,15 @@ from mcp.types import (
 @pytest.mark.anyio
 async def test_completion_handler_receives_context():
     """Test that the completion handler receives context correctly."""
-    server = Server("test-server")
-
     # Track what the handler receives
-    received_args: dict[str, Any] = {}
+    received_params: CompleteRequestParams | None = None
 
-    @server.completion()
-    async def handle_completion(
-        ref: PromptReference | ResourceTemplateReference,
-        argument: CompletionArgument,
-        context: CompletionContext | None,
-    ) -> Completion | None:
-        received_args["ref"] = ref
-        received_args["argument"] = argument
-        received_args["context"] = context
+    async def handle_completion(ctx: ServerRequestContext, params: CompleteRequestParams) -> CompleteResult:
+        nonlocal received_params
+        received_params = params
+        return CompleteResult(completion=Completion(values=["test-completion"], total=1, has_more=False))
 
-        # Return test completion
-        return Completion(values=["test-completion"], total=1, has_more=False)
+    server = Server("test-server", on_completion=handle_completion)
 
     async with Client(server) as client:
         # Test with context
@@ -45,28 +35,23 @@ async def test_completion_handler_receives_context():
         )
 
         # Verify handler received the context
-        assert received_args["context"] is not None
-        assert received_args["context"].arguments == {"previous": "value"}
+        assert received_params is not None
+        assert received_params.context is not None
+        assert received_params.context.arguments == {"previous": "value"}
         assert result.completion.values == ["test-completion"]
 
 
 @pytest.mark.anyio
 async def test_completion_backward_compatibility():
     """Test that completion works without context (backward compatibility)."""
-    server = Server("test-server")
-
     context_was_none = False
 
-    @server.completion()
-    async def handle_completion(
-        ref: PromptReference | ResourceTemplateReference,
-        argument: CompletionArgument,
-        context: CompletionContext | None,
-    ) -> Completion | None:
+    async def handle_completion(ctx: ServerRequestContext, params: CompleteRequestParams) -> CompleteResult:
         nonlocal context_was_none
-        context_was_none = context is None
+        context_was_none = params.context is None
+        return CompleteResult(completion=Completion(values=["no-context-completion"], total=1, has_more=False))
 
-        return Completion(values=["no-context-completion"], total=1, has_more=False)
+    server = Server("test-server", on_completion=handle_completion)
 
     async with Client(server) as client:
         # Test without context
@@ -82,30 +67,31 @@ async def test_completion_backward_compatibility():
 @pytest.mark.anyio
 async def test_dependent_completion_scenario():
     """Test a real-world scenario with dependent completions."""
-    server = Server("test-server")
 
-    @server.completion()
-    async def handle_completion(
-        ref: PromptReference | ResourceTemplateReference,
-        argument: CompletionArgument,
-        context: CompletionContext | None,
-    ) -> Completion | None:
+    async def handle_completion(ctx: ServerRequestContext, params: CompleteRequestParams) -> CompleteResult:
         # Simulate database/table completion scenario
-        if isinstance(ref, ResourceTemplateReference):
-            if ref.uri == "db://{database}/{table}":
-                if argument.name == "database":
-                    # Complete database names
-                    return Completion(values=["users_db", "products_db", "analytics_db"], total=3, has_more=False)
-                elif argument.name == "table":
-                    # Complete table names based on selected database
-                    if context and context.arguments:
-                        db = context.arguments.get("database")
-                        if db == "users_db":
-                            return Completion(values=["users", "sessions", "permissions"], total=3, has_more=False)
-                        elif db == "products_db":
-                            return Completion(values=["products", "categories", "inventory"], total=3, has_more=False)
+        assert isinstance(params.ref, ResourceTemplateReference)
+        assert params.ref.uri == "db://{database}/{table}"
 
-        return Completion(values=[], total=0, has_more=False)  # pragma: no cover
+        if params.argument.name == "database":
+            return CompleteResult(
+                completion=Completion(values=["users_db", "products_db", "analytics_db"], total=3, has_more=False)
+            )
+
+        assert params.argument.name == "table"
+        assert params.context and params.context.arguments
+        db = params.context.arguments.get("database")
+        if db == "users_db":
+            return CompleteResult(
+                completion=Completion(values=["users", "sessions", "permissions"], total=3, has_more=False)
+            )
+        else:
+            assert db == "products_db"
+            return CompleteResult(
+                completion=Completion(values=["products", "categories", "inventory"], total=3, has_more=False)
+            )
+
+    server = Server("test-server", on_completion=handle_completion)
 
     async with Client(server) as client:
         # First, complete database
@@ -136,27 +122,20 @@ async def test_dependent_completion_scenario():
 @pytest.mark.anyio
 async def test_completion_error_on_missing_context():
     """Test that server can raise error when required context is missing."""
-    server = Server("test-server")
 
-    @server.completion()
-    async def handle_completion(
-        ref: PromptReference | ResourceTemplateReference,
-        argument: CompletionArgument,
-        context: CompletionContext | None,
-    ) -> Completion | None:
-        if isinstance(ref, ResourceTemplateReference):
-            if ref.uri == "db://{database}/{table}":
-                if argument.name == "table":
-                    # Check if database context is provided
-                    if not context or not context.arguments or "database" not in context.arguments:
-                        # Raise an error instead of returning error as completion
-                        raise ValueError("Please select a database first to see available tables")
-                    # Normal completion if context is provided
-                    db = context.arguments.get("database")
-                    if db == "test_db":
-                        return Completion(values=["users", "orders", "products"], total=3, has_more=False)
+    async def handle_completion(ctx: ServerRequestContext, params: CompleteRequestParams) -> CompleteResult:
+        assert isinstance(params.ref, ResourceTemplateReference)
+        assert params.ref.uri == "db://{database}/{table}"
+        assert params.argument.name == "table"
 
-        return Completion(values=[], total=0, has_more=False)  # pragma: no cover
+        if not params.context or not params.context.arguments or "database" not in params.context.arguments:
+            raise ValueError("Please select a database first to see available tables")
+
+        db = params.context.arguments.get("database")
+        assert db == "test_db"
+        return CompleteResult(completion=Completion(values=["users", "orders", "products"], total=3, has_more=False))
+
+    server = Server("test-server", on_completion=handle_completion)
 
     async with Client(server) as client:
         # Try to complete table without database context - should raise error