diff --git a/tests/tools/test_tool.py b/tests/tools/test_tool.py index 24b10ce1f..67ebe9da7 100644 --- a/tests/tools/test_tool.py +++ b/tests/tools/test_tool.py @@ -1,6 +1,6 @@ import json from dataclasses import dataclass -from typing import Annotated, Any, TypedDict +from typing import Annotated, Any import pytest from mcp.types import ( @@ -11,6 +11,7 @@ from mcp.types import ( TextResourceContents, ) from pydantic import AnyUrl, BaseModel, Field, TypeAdapter +from typing_extensions import TypedDict from fastmcp.tools.tool import Tool, _convert_to_content from fastmcp.utilities.types import Audio, File, Image diff --git a/tests/tools/test_tool_transform.py b/tests/tools/test_tool_transform.py index 9a31db81d..5e976b381 100644 --- a/tests/tools/test_tool_transform.py +++ b/tests/tools/test_tool_transform.py @@ -254,7 +254,7 @@ async def test_forward_raw_without_argument_mapping(add_tool): async def test_custom_fn_with_kwargs_and_no_transform_args(add_tool): - async def custom_fn(extra: int, **kwargs) -> ToolResult: + async def custom_fn(extra: int, **kwargs) -> int: sum = await forward(**kwargs) return int(sum.content[0].text) + extra # type: ignore[attr-defined] @@ -1216,10 +1216,10 @@ class TestTransformToolOutputSchema: ): """Test transformation handling of mixed content types.""" - async def custom_fn(x: int) -> ToolResult: + async def custom_fn(x: int): # Return mixed content including ToolResult if x == 1: - return ["text", {"data": x}] # type: ignore[return-value] + return ["text", {"data": x}] else: # Return ToolResult directly return ToolResult(