diff --git a/tests/server/test_session.py b/tests/server/test_session.py
index d00eda875..154c3a368 100644
--- a/tests/server/test_session.py
+++ b/tests/server/test_session.py
@@ -13,12 +13,13 @@
     ClientNotification,
     Completion,
     CompletionArgument,
+    CompletionContext,
     CompletionsCapability,
     InitializedNotification,
     PromptReference,
     PromptsCapability,
-    ResourceReference,
     ResourcesCapability,
+    ResourceTemplateReference,
     ServerCapabilities,
 )
 
@@ -109,7 +110,11 @@ async def list_resources():
 
     # Add a complete handler
     @server.completion()
-    async def complete(ref: PromptReference | ResourceReference, argument: CompletionArgument):
+    async def complete(
+        ref: PromptReference | ResourceTemplateReference,
+        argument: CompletionArgument,
+        context: CompletionContext | None,
+    ) -> Completion | None:
         return Completion(
             values=["completion1", "completion2"],
         )