Add regression tests for functools.wraps + Context (#2524) (#2566)

Follow-up to PR #2563 which fixed the signature handling in
create_function_without_params. These tests ensure the fix
works end-to-end for all object types that support Context injection.

tests/prompts/test_prompt_manager.py

@@ -1,8 +1,9 @@
+import functools
 from typing import Annotated
 
 import pytest
 
-from fastmcp import Context
+from fastmcp import Context, FastMCP
 from fastmcp.exceptions import NotFoundError, PromptError
 from fastmcp.prompts import Prompt
 from fastmcp.prompts.prompt import FunctionPrompt, PromptMessage, TextContent
@@ -434,3 +435,30 @@ class TestContextHandling:
             return str(x)
 
         Prompt.from_function(prompt_with_context)
+
+    async def test_context_with_functools_wraps_decorator(self):
+        """Regression test for #2524: decorated prompts with Context should work."""
+
+        def custom_decorator(func):
+            @functools.wraps(func)
+            async def wrapper(*args, **kwargs):
+                return await func(*args, **kwargs)
+
+            return wrapper
+
+        @custom_decorator
+        async def decorated_prompt(ctx: Context, topic: str) -> str:
+            assert isinstance(ctx, Context)
+            return f"Write about {topic}"
+
+        prompt = Prompt.from_function(decorated_prompt)
+
+        # Verify ctx is excluded from arguments
+        assert "ctx" not in [arg.name for arg in prompt.arguments or []]
+
+        mcp = FastMCP()
+        context = Context(fastmcp=mcp)
+
+        async with context:
+            messages = await prompt.render(arguments={"topic": "cats"})
+            assert messages[0].content.text == "Write about cats"  # type: ignore[attr-defined]

tests/resources/test_resource_template.py

@@ -1,10 +1,11 @@
+import functools
 import json
 from urllib.parse import quote
 
 import pytest
 from pydantic import BaseModel
 
-from fastmcp import Context
+from fastmcp import Context, FastMCP
 from fastmcp.resources import ResourceTemplate
 from fastmcp.resources.resource import FunctionResource
 from fastmcp.resources.template import match_uri_template
@@ -693,8 +694,6 @@ class TestContextHandling:
         )
 
         # Even for optional context, we need to provide a context
-        from fastmcp import FastMCP
-
         mcp = FastMCP()
         context = Context(fastmcp=mcp)
 
@@ -708,6 +707,35 @@ class TestContextHandling:
             content = await resource.read()
             assert content == "42"
 
+    async def test_context_with_functools_wraps_decorator(self):
+        """Regression test for #2524: decorated templates with Context should work."""
+
+        def custom_decorator(func):
+            @functools.wraps(func)
+            async def wrapper(*args, **kwargs):
+                return await func(*args, **kwargs)
+
+            return wrapper
+
+        @custom_decorator
+        async def decorated_template(ctx: Context, item_id: int) -> str:
+            assert isinstance(ctx, Context)
+            return f"item: {item_id}"
+
+        template = ResourceTemplate.from_function(
+            fn=decorated_template,
+            uri_template="test://{item_id}",
+            name="test",
+        )
+
+        mcp = FastMCP()
+        context = Context(fastmcp=mcp)
+
+        async with context:
+            resource = await template.create_resource("test://42", {"item_id": 42})
+            content = await resource.read()
+            assert content == "item: 42"
+
 
 class TestQueryParameterExtraction:
     """Test basic query parameter extraction from URIs."""

tests/server/test_server_interactions.py

@@ -1,5 +1,6 @@
 import base64
 import datetime
+import functools
 import json
 import uuid
 from dataclasses import dataclass
@@ -1310,6 +1311,34 @@ class TestToolContextInjection:
             result = await client.call_tool("MyTool", {"x": 2})
             assert result.data == 3
 
+    async def test_decorated_tool_with_functools_wraps(self):
+        """Regression test for #2524: @mcp.tool with functools.wraps decorator."""
+
+        def custom_decorator(func):
+            @functools.wraps(func)
+            async def wrapper(*args, **kwargs):
+                return await func(*args, **kwargs)
+
+            return wrapper
+
+        mcp = FastMCP()
+
+        @mcp.tool
+        @custom_decorator
+        async def decorated_tool(ctx: Context, query: str) -> str:
+            assert isinstance(ctx, Context)
+            return f"query: {query}"
+
+        async with Client(mcp) as client:
+            # Verify ctx is not in the schema
+            tools = await client.list_tools()
+            tool = next(t for t in tools if t.name == "decorated_tool")
+            assert "ctx" not in tool.inputSchema.get("properties", {})
+
+            # Verify the tool works
+            result = await client.call_tool("decorated_tool", {"query": "test"})
+            assert result.data == "query: test"
+
 
 class TestToolEnabled:
     async def test_toggle_enabled(self):

tests/tools/test_tool_manager.py

@@ -1,3 +1,4 @@
+import functools
 import json
 import logging
 import uuid
@@ -816,6 +817,35 @@ class TestContextHandling:
             ):
                 await manager.call_tool("tool_with_context", {"x": 42})
 
+    async def test_context_with_functools_wraps_decorator(self):
+        """Regression test for #2524: decorated tools with Context should work."""
+
+        def custom_decorator(func):
+            @functools.wraps(func)
+            async def wrapper(*args, **kwargs):
+                return await func(*args, **kwargs)
+
+            return wrapper
+
+        @custom_decorator
+        async def decorated_tool(ctx: Context, query: str) -> str:
+            assert isinstance(ctx, Context)
+            return f"query: {query}"
+
+        manager = ToolManager()
+        tool = Tool.from_function(decorated_tool)
+        manager.add_tool(tool)
+
+        # Verify ctx is excluded from schema
+        assert "ctx" not in json.dumps(tool.parameters)
+
+        mcp = FastMCP()
+        context = Context(fastmcp=mcp)
+
+        async with context:
+            result = await manager.call_tool("decorated_tool", {"query": "test"})
+            assert result.structured_content == {"result": "query: test"}
+
 
 class TestCustomToolNames:
     """Test adding tools with custom names that differ from their function names."""