mirror of
https://github.com/PrefectHQ/fastmcp.git
synced 2026-08-19 20:14:17 +02:00
266 lines
9.1 KiB
Python
266 lines
9.1 KiB
Python
import pytest
|
|
from mcp.types import ImageContent, TextContent
|
|
from pydantic import BaseModel
|
|
|
|
from fastmcp import Image
|
|
from fastmcp.tools.tool import Tool
|
|
|
|
|
|
class TestToolFromFunction:
|
|
def test_basic_function(self):
|
|
"""Test registering and running a basic function."""
|
|
|
|
def add(a: int, b: int) -> int:
|
|
"""Add two numbers."""
|
|
return a + b
|
|
|
|
tool = Tool.from_function(add)
|
|
|
|
assert tool.name == "add"
|
|
assert tool.description == "Add two numbers."
|
|
assert tool.parameters["properties"]["a"]["type"] == "integer"
|
|
assert tool.parameters["properties"]["b"]["type"] == "integer"
|
|
|
|
async def test_async_function(self):
|
|
"""Test registering and running an async function."""
|
|
|
|
async def fetch_data(url: str) -> str:
|
|
"""Fetch data from URL."""
|
|
return f"Data from {url}"
|
|
|
|
tool = Tool.from_function(fetch_data)
|
|
|
|
assert tool.name == "fetch_data"
|
|
assert tool.description == "Fetch data from URL."
|
|
assert tool.parameters["properties"]["url"]["type"] == "string"
|
|
|
|
def test_pydantic_model_function(self):
|
|
"""Test registering a function that takes a Pydantic model."""
|
|
|
|
class UserInput(BaseModel):
|
|
name: str
|
|
age: int
|
|
|
|
def create_user(user: UserInput, flag: bool) -> dict:
|
|
"""Create a new user."""
|
|
return {"id": 1, **user.model_dump()}
|
|
|
|
tool = Tool.from_function(create_user)
|
|
|
|
assert tool.name == "create_user"
|
|
assert tool.description == "Create a new user."
|
|
assert "name" in tool.parameters["$defs"]["UserInput"]["properties"]
|
|
assert "age" in tool.parameters["$defs"]["UserInput"]["properties"]
|
|
assert "flag" in tool.parameters["properties"]
|
|
|
|
async def test_tool_with_image_return(self):
|
|
def image_tool(data: bytes) -> Image:
|
|
return Image(data=data)
|
|
|
|
tool = Tool.from_function(image_tool)
|
|
|
|
result = await tool.run({"data": "test.png"})
|
|
assert tool.parameters["properties"]["data"]["type"] == "string"
|
|
assert isinstance(result[0], ImageContent)
|
|
|
|
def test_non_callable_fn(self):
|
|
with pytest.raises(TypeError, match="not a callable object"):
|
|
Tool.from_function(1) # type: ignore
|
|
|
|
def test_lambda(self):
|
|
tool = Tool.from_function(lambda x: x, name="my_tool")
|
|
assert tool.name == "my_tool"
|
|
|
|
def test_lambda_with_no_name(self):
|
|
with pytest.raises(
|
|
ValueError, match="You must provide a name for lambda functions"
|
|
):
|
|
Tool.from_function(lambda x: x)
|
|
|
|
def test_private_arguments(self):
|
|
def add(_a: int, _b: int) -> int:
|
|
"""Add two numbers."""
|
|
return _a + _b
|
|
|
|
tool = Tool.from_function(add)
|
|
assert tool.parameters["properties"]["_a"]["type"] == "integer"
|
|
assert tool.parameters["properties"]["_b"]["type"] == "integer"
|
|
|
|
def test_tool_with_varargs_not_allowed(self):
|
|
def func(a: int, b: int, *args: int) -> int:
|
|
"""Add two numbers."""
|
|
return a + b
|
|
|
|
with pytest.raises(
|
|
ValueError, match=r"Functions with \*args are not supported as tools"
|
|
):
|
|
Tool.from_function(func)
|
|
|
|
def test_tool_with_varkwargs_not_allowed(self):
|
|
def func(a: int, b: int, **kwargs: int) -> int:
|
|
"""Add two numbers."""
|
|
return a + b
|
|
|
|
with pytest.raises(
|
|
ValueError, match=r"Functions with \*\*kwargs are not supported as tools"
|
|
):
|
|
Tool.from_function(func)
|
|
|
|
async def test_instance_method(self):
|
|
class MyClass:
|
|
def add(self, x: int, y: int) -> int:
|
|
"""Add two numbers."""
|
|
return x + y
|
|
|
|
obj = MyClass()
|
|
|
|
tool = Tool.from_function(obj.add)
|
|
assert tool.name == "add"
|
|
assert tool.description == "Add two numbers."
|
|
assert "self" not in tool.parameters["properties"]
|
|
|
|
async def test_instance_method_with_varargs_not_allowed(self):
|
|
class MyClass:
|
|
def add(self, x: int, y: int, *args: int) -> int:
|
|
"""Add two numbers."""
|
|
return x + y
|
|
|
|
obj = MyClass()
|
|
|
|
with pytest.raises(
|
|
ValueError, match=r"Functions with \*args are not supported as tools"
|
|
):
|
|
Tool.from_function(obj.add)
|
|
|
|
async def test_instance_method_with_varkwargs_not_allowed(self):
|
|
class MyClass:
|
|
def add(self, x: int, y: int, **kwargs: int) -> int:
|
|
"""Add two numbers."""
|
|
return x + y
|
|
|
|
obj = MyClass()
|
|
|
|
with pytest.raises(
|
|
ValueError, match=r"Functions with \*\*kwargs are not supported as tools"
|
|
):
|
|
Tool.from_function(obj.add)
|
|
|
|
async def test_classmethod(self):
|
|
class MyClass:
|
|
x: int = 10
|
|
|
|
|
|
class TestToolJsonParsing:
|
|
"""Tests for Tool's JSON pre-parsing functionality."""
|
|
|
|
async def test_json_string_arguments(self):
|
|
"""Test that JSON string arguments are parsed and validated correctly"""
|
|
|
|
def simple_func(x: int, y: list[str]) -> str:
|
|
return f"{x}-{','.join(y)}"
|
|
|
|
# Create a tool to use its JSON pre-parsing logic
|
|
tool = Tool.from_function(simple_func)
|
|
|
|
# Prepare arguments where some are JSON strings
|
|
json_args = {
|
|
"x": 1,
|
|
"y": '["a", "b", "c"]', # JSON string
|
|
}
|
|
|
|
# Run the tool which will do JSON parsing
|
|
result = await tool.run(json_args)
|
|
assert len(result) == 1
|
|
assert isinstance(result[0], TextContent)
|
|
assert result[0].text == "1-a,b,c"
|
|
|
|
async def test_str_vs_list_str(self):
|
|
"""Test handling of string vs list[str] type annotations."""
|
|
|
|
def func_with_str_types(str_or_list: str | list[str]) -> str | list[str]:
|
|
return str_or_list
|
|
|
|
tool = Tool.from_function(func_with_str_types)
|
|
|
|
# Test regular string input (should remain a string)
|
|
result = await tool.run({"str_or_list": "hello"})
|
|
assert len(result) == 1
|
|
assert isinstance(result[0], TextContent)
|
|
assert result[0].text == "hello"
|
|
|
|
# Test JSON string input (should be parsed as a string)
|
|
result = await tool.run({"str_or_list": '"hello"'})
|
|
assert len(result) == 1
|
|
assert isinstance(result[0], TextContent)
|
|
assert result[0].text == "hello"
|
|
|
|
# Test JSON list input (should be parsed as a list)
|
|
result = await tool.run({"str_or_list": '["hello", "world"]'})
|
|
assert len(result) == 1
|
|
assert isinstance(result[0], TextContent)
|
|
|
|
# The exact formatting might vary, so we just check that it contains the key elements
|
|
text_without_whitespace = result[0].text.replace(" ", "").replace("\n", "")
|
|
assert "hello" in text_without_whitespace
|
|
assert "world" in text_without_whitespace
|
|
assert "[" in text_without_whitespace
|
|
assert "]" in text_without_whitespace
|
|
|
|
async def test_keep_str_as_str(self):
|
|
"""Test that string arguments are kept as strings when they're not valid JSON"""
|
|
|
|
def func_with_str_types(string: str) -> str:
|
|
return string
|
|
|
|
tool = Tool.from_function(func_with_str_types)
|
|
|
|
# Invalid JSON should remain a string
|
|
invalid_json = "{'nice to meet you': 'hello', 'goodbye': 5}"
|
|
result = await tool.run({"string": invalid_json})
|
|
assert len(result) == 1
|
|
assert isinstance(result[0], TextContent)
|
|
assert result[0].text == invalid_json
|
|
|
|
async def test_keep_str_union_as_str(self):
|
|
"""Test that string arguments are kept as strings when parsing would create an invalid value"""
|
|
|
|
def func_with_str_types(
|
|
string: str | dict[int, str] | None,
|
|
) -> str | dict[int, str] | None:
|
|
return string
|
|
|
|
tool = Tool.from_function(func_with_str_types)
|
|
|
|
# Invalid JSON for the union type should remain a string
|
|
invalid_json = "{'nice to meet you': 'hello', 'goodbye': 5}"
|
|
result = await tool.run({"string": invalid_json})
|
|
assert len(result) == 1
|
|
assert isinstance(result[0], TextContent)
|
|
assert result[0].text == invalid_json
|
|
|
|
async def test_complex_type_validation(self):
|
|
"""Test that parsed JSON is validated against complex types"""
|
|
|
|
class SomeModel(BaseModel):
|
|
x: int
|
|
y: dict[int, str]
|
|
|
|
def func_with_complex_type(data: SomeModel) -> SomeModel:
|
|
return data
|
|
|
|
tool = Tool.from_function(func_with_complex_type)
|
|
|
|
# Valid JSON for the model
|
|
valid_json = '{"x": 1, "y": {"1": "hello"}}'
|
|
result = await tool.run({"data": valid_json})
|
|
assert len(result) == 1
|
|
assert isinstance(result[0], TextContent)
|
|
assert '"x": 1' in result[0].text
|
|
assert '"y": {' in result[0].text
|
|
assert '"1": "hello"' in result[0].text
|
|
|
|
# Invalid JSON for the model (y has string keys, not int keys)
|
|
# Should throw a validation error
|
|
invalid_json = '{"x": 1, "y": {"invalid": "hello"}}'
|
|
with pytest.raises(Exception):
|
|
await tool.run({"data": invalid_json})
|