From 0cf12fa30cfea31afb64a95608d4627231ab86d6 Mon Sep 17 00:00:00 2001 From: Jeremiah Lowin <153965+jlowin@users.noreply.github.com> Date: Sat, 6 Dec 2025 12:32:39 -0500 Subject: [PATCH] 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 | 30 +++++++++++++++++++- tests/resources/test_resource_template.py | 34 +++++++++++++++++++++-- tests/server/test_server_interactions.py | 29 +++++++++++++++++++ tests/tools/test_tool_manager.py | 30 ++++++++++++++++++++ 4 files changed, 119 insertions(+), 4 deletions(-) diff --git a/tests/prompts/test_prompt_manager.py b/tests/prompts/test_prompt_manager.py index 61b0e991e..076c27279 100644 --- a/tests/prompts/test_prompt_manager.py +++ b/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] diff --git a/tests/resources/test_resource_template.py b/tests/resources/test_resource_template.py index d7a2f6154..ea6cf0e43 100644 --- a/tests/resources/test_resource_template.py +++ b/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.""" diff --git a/tests/server/test_server_interactions.py b/tests/server/test_server_interactions.py index 238893151..8eb2d2769 100644 --- a/tests/server/test_server_interactions.py +++ b/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): diff --git a/tests/tools/test_tool_manager.py b/tests/tools/test_tool_manager.py index 1ece0dfd2..128bdea79 100644 --- a/tests/tools/test_tool_manager.py +++ b/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."""