Jeremiah Lowin commited on
Commit
a004953
·
1 Parent(s): 469d590

Add tool tests

Browse files
Files changed (1) hide show
  1. tests/test_tool_manager.py +119 -0
tests/test_tool_manager.py ADDED
@@ -0,0 +1,119 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """Test tool registration and execution."""
2
+
3
+ import pytest
4
+ from pydantic import BaseModel
5
+
6
+ from fastmcp.exceptions import ToolError
7
+ from fastmcp.tools import ToolManager
8
+
9
+
10
+ class TestAddTools:
11
+ def test_basic_function(self):
12
+ """Test registering and running a basic function."""
13
+
14
+ def add(a: int, b: int) -> int:
15
+ """Add two numbers."""
16
+ return a + b
17
+
18
+ manager = ToolManager()
19
+ manager.add_tool(add)
20
+
21
+ tool = manager.get_tool("add")
22
+ assert tool is not None
23
+ assert tool.name == "add"
24
+ assert tool.description == "Add two numbers."
25
+ assert tool.is_async is False
26
+ assert tool.parameters["properties"]["a"]["type"] == "integer"
27
+ assert tool.parameters["properties"]["b"]["type"] == "integer"
28
+
29
+ async def test_async_function(self):
30
+ """Test registering and running an async function."""
31
+
32
+ async def fetch_data(url: str) -> str:
33
+ """Fetch data from URL."""
34
+ return f"Data from {url}"
35
+
36
+ manager = ToolManager()
37
+ manager.add_tool(fetch_data)
38
+
39
+ tool = manager.get_tool("fetch_data")
40
+ assert tool is not None
41
+ assert tool.name == "fetch_data"
42
+ assert tool.description == "Fetch data from URL."
43
+ assert tool.is_async is True
44
+ assert tool.parameters["properties"]["url"]["type"] == "string"
45
+
46
+ def test_pydantic_model_function(self):
47
+ """Test registering a function that takes a Pydantic model."""
48
+
49
+ class UserInput(BaseModel):
50
+ name: str
51
+ age: int
52
+
53
+ def create_user(user: UserInput, flag: bool) -> dict:
54
+ """Create a new user."""
55
+ return {"id": 1, **user.model_dump()}
56
+
57
+ manager = ToolManager()
58
+ manager.add_tool(create_user)
59
+
60
+ tool = manager.get_tool("create_user")
61
+ assert tool is not None
62
+ assert tool.name == "create_user"
63
+ assert tool.description == "Create a new user."
64
+ assert tool.is_async is False
65
+ assert "name" in tool.parameters["$defs"]["UserInput"]["properties"]
66
+ assert "age" in tool.parameters["$defs"]["UserInput"]["properties"]
67
+ assert "flag" in tool.parameters["properties"]
68
+
69
+ def test_add_invalid_tool(self):
70
+ manager = ToolManager()
71
+ with pytest.raises(AttributeError):
72
+ manager.add_tool(1)
73
+
74
+
75
+ class TestCallTools:
76
+ async def test_call_tool(self):
77
+ def add(a: int, b: int) -> int:
78
+ """Add two numbers."""
79
+ return a + b
80
+
81
+ manager = ToolManager()
82
+ manager.add_tool(add)
83
+ result = await manager.call_tool("add", {"a": 1, "b": 2})
84
+ assert result == 3
85
+
86
+ async def test_call_async_tool(self):
87
+ async def double(n: int) -> int:
88
+ """Double a number."""
89
+ return n * 2
90
+
91
+ manager = ToolManager()
92
+ manager.add_tool(double)
93
+ result = await manager.call_tool("double", {"n": 5})
94
+ assert result == 10
95
+
96
+ async def test_call_tool_with_default_args(self):
97
+ def add(a: int, b: int = 1) -> int:
98
+ """Add two numbers."""
99
+ return a + b
100
+
101
+ manager = ToolManager()
102
+ manager.add_tool(add)
103
+ result = await manager.call_tool("add", {"a": 1})
104
+ assert result == 2
105
+
106
+ async def test_call_tool_with_missing_args(self):
107
+ def add(a: int, b: int) -> int:
108
+ """Add two numbers."""
109
+ return a + b
110
+
111
+ manager = ToolManager()
112
+ manager.add_tool(add)
113
+ with pytest.raises(ToolError):
114
+ await manager.call_tool("add", {"a": 1})
115
+
116
+ async def test_call_unknown_tool(self):
117
+ manager = ToolManager()
118
+ with pytest.raises(ToolError):
119
+ await manager.call_tool("unknown", {"a": 1})