248 lines
7.4 KiB
Python
248 lines
7.4 KiB
Python
import pytest
|
|
import inspect
|
|
from typing import Dict, Any, List, Optional, Callable
|
|
|
|
from ai.mcp_server import DummyMCPServer
|
|
|
|
|
|
@pytest.fixture
|
|
def mcp_server(session):
|
|
return DummyMCPServer(session, None)
|
|
|
|
|
|
def test_register_tool_basic(mcp_server):
|
|
"""Test basic tool registration"""
|
|
|
|
# Define a simple test function
|
|
def test_func(param1: str, param2: int):
|
|
return f"Test {param1} {param2}"
|
|
|
|
# Register the tool
|
|
result = mcp_server.register_tool("test_tool", test_func)
|
|
|
|
# Verify the tool was registered
|
|
assert "test_tool" in mcp_server.available_tools
|
|
assert mcp_server.available_tools["test_tool"]["name"] == "test_tool"
|
|
assert mcp_server.available_tools["test_tool"]["handler"] == test_func
|
|
|
|
# Verify method chaining works
|
|
assert result is mcp_server
|
|
|
|
|
|
def test_register_tool_with_description(mcp_server):
|
|
"""Test tool registration with a custom description"""
|
|
|
|
def test_func():
|
|
return "test"
|
|
|
|
mcp_server.register_tool("test_tool", test_func, description="Custom description")
|
|
|
|
assert mcp_server.available_tools["test_tool"]["description"] == "Custom description"
|
|
|
|
|
|
def test_register_tool_without_description(mcp_server):
|
|
"""Test tool registration without a description"""
|
|
|
|
def test_func():
|
|
return "test"
|
|
|
|
mcp_server.register_tool("test_tool", test_func)
|
|
|
|
assert mcp_server.available_tools["test_tool"]["description"] == "Tool test_tool"
|
|
|
|
|
|
def test_register_tool_parameter_types(mcp_server):
|
|
"""Test parameter type inference"""
|
|
|
|
def test_func(str_param: str, int_param: int, float_param: float, bool_param: bool, untyped_param):
|
|
return "test"
|
|
|
|
mcp_server.register_tool("test_tool", test_func)
|
|
|
|
params = mcp_server.available_tools["test_tool"]["parameters"]
|
|
|
|
assert params["str_param"]["type"] == "string"
|
|
assert params["int_param"]["type"] == "integer"
|
|
assert params["float_param"]["type"] == "number"
|
|
assert params["bool_param"]["type"] == "boolean"
|
|
assert params["untyped_param"]["type"] == "string" # Default type for untyped parameters
|
|
|
|
|
|
def test_register_tool_parameter_descriptions(mcp_server):
|
|
"""Test parameter descriptions"""
|
|
|
|
def test_func(param1, param2):
|
|
return "test"
|
|
|
|
mcp_server.register_tool("test_tool", test_func)
|
|
|
|
params = mcp_server.available_tools["test_tool"]["parameters"]
|
|
|
|
assert params["param1"]["description"] == "Parameter param1"
|
|
assert params["param2"]["description"] == "Parameter param2"
|
|
|
|
|
|
def test_register_tool_with_sphinx_docstring(mcp_server):
|
|
"""Test parameter descriptions from Sphinx-style docstrings"""
|
|
|
|
def test_func(name: str, age: int):
|
|
"""Test function with Sphinx docstring
|
|
|
|
:param name: The person's name
|
|
:param age: The person's age in years
|
|
:return: A greeting message
|
|
"""
|
|
return f"Hello {name}, you are {age} years old!"
|
|
|
|
mcp_server.register_tool("sphinx_doc_tool", test_func)
|
|
|
|
params = mcp_server.available_tools["sphinx_doc_tool"]["parameters"]
|
|
|
|
assert params["name"]["description"] == "The person's name"
|
|
assert params["age"]["description"] == "The person's age in years"
|
|
|
|
|
|
def test_register_tool_with_google_docstring(mcp_server):
|
|
"""Test parameter descriptions from Google-style docstrings"""
|
|
|
|
def test_func(name: str, age: int, height: float):
|
|
"""Test function with Google-style docstring
|
|
|
|
Args:
|
|
name: The person's name
|
|
age: The person's age in years
|
|
height: The person's height in meters
|
|
|
|
Returns:
|
|
A greeting message
|
|
"""
|
|
return f"Hello {name}, you are {age} years old and {height}m tall!"
|
|
|
|
mcp_server.register_tool("google_doc_tool", test_func)
|
|
|
|
params = mcp_server.available_tools["google_doc_tool"]["parameters"]
|
|
|
|
assert params["name"]["description"] == "The person's name"
|
|
assert params["age"]["description"] == "The person's age in years"
|
|
assert params["height"]["description"] == "The person's height in meters"
|
|
|
|
|
|
def test_register_tool_with_parameters_keyword(mcp_server):
|
|
"""Test parameter descriptions with 'Parameters:' keyword instead of 'Args:'"""
|
|
|
|
def test_func(x: int, y: int):
|
|
"""Test function with Parameters keyword
|
|
|
|
Parameters:
|
|
x: The x coordinate
|
|
y: The y coordinate
|
|
|
|
Returns:
|
|
The sum of coordinates
|
|
"""
|
|
return x + y
|
|
|
|
mcp_server.register_tool("parameters_doc_tool", test_func)
|
|
|
|
params = mcp_server.available_tools["parameters_doc_tool"]["parameters"]
|
|
|
|
assert params["x"]["description"] == "The x coordinate"
|
|
assert params["y"]["description"] == "The y coordinate"
|
|
|
|
|
|
def test_register_tool_with_mixed_docstrings(mcp_server):
|
|
"""Test parameter descriptions with mixed docstring styles"""
|
|
|
|
def test_func(a: int, b: str, c: float):
|
|
"""Test function with mixed docstring styles
|
|
|
|
:param a: Parameter a from Sphinx style
|
|
|
|
Args:
|
|
b: Parameter b from Google style
|
|
c: Parameter c from Google style
|
|
"""
|
|
return f"{a} {b} {c}"
|
|
|
|
mcp_server.register_tool("mixed_doc_tool", test_func)
|
|
|
|
params = mcp_server.available_tools["mixed_doc_tool"]["parameters"]
|
|
|
|
assert params["a"]["description"] == "Parameter a from Sphinx style"
|
|
assert params["b"]["description"] == "Parameter b from Google style"
|
|
assert params["c"]["description"] == "Parameter c from Google style"
|
|
|
|
|
|
def test_register_tool_with_missing_docstrings(mcp_server):
|
|
"""Test parameter descriptions when some parameters are missing from docstring"""
|
|
|
|
def test_func(a: int, b: str, c: float):
|
|
"""Test function with incomplete docstring
|
|
|
|
Args:
|
|
a: Parameter a description
|
|
"""
|
|
return f"{a} {b} {c}"
|
|
|
|
mcp_server.register_tool("incomplete_doc_tool", test_func)
|
|
|
|
params = mcp_server.available_tools["incomplete_doc_tool"]["parameters"]
|
|
|
|
assert params["a"]["description"] == "Parameter a description"
|
|
assert params["b"]["description"] == "Parameter b" # Default description
|
|
assert params["c"]["description"] == "Parameter c" # Default description
|
|
|
|
|
|
async def test_tool_can_be_called(mcp_server):
|
|
"""Test that a registered tool can be called through call_tool method"""
|
|
|
|
# Define a simple test function
|
|
def test_func(value: int):
|
|
return value * 2
|
|
|
|
# Register the tool
|
|
mcp_server.register_tool("multiply", test_func)
|
|
|
|
# Call the tool
|
|
result = await mcp_server.call_tool("multiply", {"value": 5})
|
|
|
|
# Verify the result
|
|
assert result["success"] is True
|
|
assert result["result"] == 10
|
|
assert result["tool_name"] == "multiply"
|
|
|
|
|
|
async def test_async_tool_can_be_called(mcp_server):
|
|
"""Test that a registered async tool can be called through call_tool method"""
|
|
|
|
# Define an async test function
|
|
async def async_test_func(value: int):
|
|
return value * 3
|
|
|
|
# Register the async tool
|
|
mcp_server.register_tool("async_multiply", async_test_func)
|
|
|
|
# Call the async tool
|
|
result = await mcp_server.call_tool("async_multiply", {"value": 5})
|
|
|
|
# Verify the result
|
|
assert result["success"] is True
|
|
assert result["result"] == 15
|
|
assert result["tool_name"] == "async_multiply"
|
|
|
|
|
|
def test_multiple_tools_registration(mcp_server):
|
|
"""Test registering multiple tools"""
|
|
|
|
def tool1(): return "tool1"
|
|
|
|
def tool2(): return "tool2"
|
|
|
|
nb_internal_tools = len(mcp_server.available_tools)
|
|
mcp_server.register_tool("tool1", tool1)
|
|
mcp_server.register_tool("tool2", tool2)
|
|
|
|
# Check both tools were registered
|
|
assert "tool1" in mcp_server.available_tools
|
|
assert "tool2" in mcp_server.available_tools
|
|
assert len(mcp_server.available_tools) == nb_internal_tools + 2 |