mirror of
https://github.com/PrefectHQ/fastmcp.git
synced 2026-08-20 20:44:17 +02:00
265 lines
9.7 KiB
Python
265 lines
9.7 KiB
Python
import json
|
|
from typing import Any
|
|
|
|
import pytest
|
|
from anyio import create_task_group
|
|
from dirty_equals import Contains
|
|
from mcp import McpError
|
|
from pydantic import AnyUrl
|
|
|
|
from fastmcp import FastMCP
|
|
from fastmcp.client import Client
|
|
from fastmcp.client.transports import FastMCPTransport
|
|
from fastmcp.exceptions import ToolError
|
|
from fastmcp.server.proxy import FastMCPProxy
|
|
|
|
USERS = [
|
|
{"id": "1", "name": "Alice", "active": True},
|
|
{"id": "2", "name": "Bob", "active": True},
|
|
{"id": "3", "name": "Charlie", "active": False},
|
|
]
|
|
|
|
|
|
@pytest.fixture
|
|
def fastmcp_server():
|
|
server = FastMCP("TestServer")
|
|
|
|
# --- Tools ---
|
|
|
|
@server.tool()
|
|
def greet(name: str) -> str:
|
|
"""Greet someone by name."""
|
|
return f"Hello, {name}!"
|
|
|
|
@server.tool()
|
|
def tool_without_description() -> str:
|
|
return "Hello?"
|
|
|
|
@server.tool()
|
|
def add(a: int, b: int) -> int:
|
|
"""Add two numbers together."""
|
|
return a + b
|
|
|
|
@server.tool()
|
|
def error_tool():
|
|
"""This tool always raises an error."""
|
|
raise ValueError("This is a test error")
|
|
|
|
# --- Resources ---
|
|
|
|
@server.resource(uri="resource://wave")
|
|
def wave() -> str:
|
|
return "👋"
|
|
|
|
@server.resource(uri="data://users")
|
|
async def get_users() -> list[dict[str, Any]]:
|
|
return USERS
|
|
|
|
@server.resource(uri="data://user/{user_id}")
|
|
async def get_user(user_id: str) -> dict[str, Any] | None:
|
|
return next((user for user in USERS if user["id"] == user_id), None)
|
|
|
|
# --- Prompts ---
|
|
|
|
@server.prompt()
|
|
def welcome(name: str) -> str:
|
|
return f"Welcome to FastMCP, {name}!"
|
|
|
|
return server
|
|
|
|
|
|
@pytest.fixture
|
|
async def proxy_server(fastmcp_server):
|
|
"""Fixture that creates a FastMCP proxy server."""
|
|
return FastMCP.as_proxy(Client(transport=FastMCPTransport(fastmcp_server)))
|
|
|
|
|
|
async def test_create_proxy(fastmcp_server):
|
|
"""Test that the proxy server properly forwards requests to the original server."""
|
|
# Create a client
|
|
client = Client(transport=FastMCPTransport(fastmcp_server))
|
|
|
|
server = FastMCPProxy.as_proxy(client)
|
|
|
|
assert isinstance(server, FastMCPProxy)
|
|
assert isinstance(server, FastMCP)
|
|
assert server.name == "FastMCP"
|
|
|
|
|
|
async def test_as_proxy_with_server(fastmcp_server):
|
|
"""FastMCP.as_proxy should accept a FastMCP instance."""
|
|
proxy = FastMCP.as_proxy(fastmcp_server)
|
|
result = await proxy._mcp_call_tool("greet", {"name": "Test"})
|
|
assert result[0].text == "Hello, Test!" # type: ignore[attr-defined]
|
|
|
|
|
|
async def test_as_proxy_with_transport(fastmcp_server):
|
|
"""FastMCP.as_proxy should accept a ClientTransport."""
|
|
proxy = FastMCP.as_proxy(FastMCPTransport(fastmcp_server))
|
|
result = await proxy._mcp_call_tool("greet", {"name": "Test"})
|
|
assert result[0].text == "Hello, Test!" # type: ignore[attr-defined]
|
|
|
|
|
|
def test_as_proxy_with_url():
|
|
"""FastMCP.as_proxy should accept a URL without connecting."""
|
|
proxy = FastMCP.as_proxy("http://example.com/mcp")
|
|
assert isinstance(proxy, FastMCPProxy)
|
|
assert repr(proxy.client.transport).startswith("<StreamableHttp(")
|
|
|
|
|
|
class TestTools:
|
|
async def test_get_tools(self, proxy_server):
|
|
tools = await proxy_server.get_tools()
|
|
assert "greet" in tools
|
|
assert "add" in tools
|
|
assert "error_tool" in tools
|
|
assert "tool_without_description" in tools
|
|
|
|
async def test_tool_without_description(self, proxy_server):
|
|
tools = await proxy_server.get_tools()
|
|
assert tools["tool_without_description"].description is None
|
|
|
|
async def test_list_tools_same_as_original(self, fastmcp_server, proxy_server):
|
|
assert (
|
|
await proxy_server._mcp_list_tools()
|
|
== await fastmcp_server._mcp_list_tools()
|
|
)
|
|
|
|
async def test_call_tool_result_same_as_original(
|
|
self, fastmcp_server: FastMCP, proxy_server: FastMCPProxy
|
|
):
|
|
result = await fastmcp_server._mcp_call_tool("greet", {"name": "Alice"})
|
|
proxy_result = await proxy_server._mcp_call_tool("greet", {"name": "Alice"})
|
|
|
|
assert result == proxy_result
|
|
|
|
async def test_call_tool_calls_tool(self, proxy_server):
|
|
async with Client(proxy_server) as client:
|
|
proxy_result = await client.call_tool("add", {"a": 1, "b": 2})
|
|
assert proxy_result[0].text == "3" # type: ignore[attr-defined]
|
|
|
|
async def test_error_tool_raises_error(self, proxy_server):
|
|
with pytest.raises(ToolError, match=""):
|
|
async with Client(proxy_server) as client:
|
|
await client.call_tool("error_tool", {})
|
|
|
|
|
|
class TestResources:
|
|
async def test_get_resources(self, proxy_server):
|
|
resources = await proxy_server.get_resources()
|
|
assert [r.uri for r in resources.values()] == Contains(
|
|
AnyUrl("data://users"),
|
|
AnyUrl("resource://wave"),
|
|
)
|
|
assert [r.name for r in resources.values()] == Contains("get_users", "wave")
|
|
|
|
async def test_list_resources_same_as_original(self, fastmcp_server, proxy_server):
|
|
assert (
|
|
await proxy_server._mcp_list_resources()
|
|
== await fastmcp_server._mcp_list_resources()
|
|
)
|
|
|
|
async def test_read_resource(self, proxy_server: FastMCPProxy):
|
|
async with Client(proxy_server) as client:
|
|
result = await client.read_resource("resource://wave")
|
|
assert result[0].text == "👋" # type: ignore[attr-defined]
|
|
|
|
async def test_read_resource_same_as_original(self, fastmcp_server, proxy_server):
|
|
async with Client(fastmcp_server) as client:
|
|
result = await client.read_resource("resource://wave")
|
|
async with Client(proxy_server) as client:
|
|
proxy_result = await client.read_resource("resource://wave")
|
|
assert proxy_result == result
|
|
|
|
async def test_read_json_resource(self, proxy_server: FastMCPProxy):
|
|
async with Client(proxy_server) as client:
|
|
result = await client.read_resource("data://users")
|
|
assert json.loads(result[0].text) == USERS # type: ignore[attr-defined]
|
|
|
|
async def test_read_resource_returns_none_if_not_found(self, proxy_server):
|
|
with pytest.raises(McpError, match="Unknown resource: resource://nonexistent"):
|
|
async with Client(proxy_server) as client:
|
|
await client.read_resource("resource://nonexistent")
|
|
|
|
|
|
class TestResourceTemplates:
|
|
async def test_get_resource_templates(self, proxy_server):
|
|
templates = await proxy_server.get_resource_templates()
|
|
assert [t.name for t in templates.values()] == Contains("get_user")
|
|
|
|
async def test_list_resource_templates_same_as_original(
|
|
self, fastmcp_server, proxy_server
|
|
):
|
|
result = await fastmcp_server._mcp_list_resource_templates()
|
|
proxy_result = await proxy_server._mcp_list_resource_templates()
|
|
assert proxy_result == result
|
|
|
|
@pytest.mark.parametrize("id", [1, 2, 3])
|
|
async def test_read_resource_template(self, proxy_server: FastMCPProxy, id: int):
|
|
async with Client(proxy_server) as client:
|
|
result = await client.read_resource(f"data://user/{id}")
|
|
assert json.loads(result[0].text) == USERS[id - 1] # type: ignore[attr-defined]
|
|
|
|
async def test_read_resource_template_same_as_original(
|
|
self, fastmcp_server, proxy_server
|
|
):
|
|
async with Client(fastmcp_server) as client:
|
|
result = await client.read_resource("data://user/1")
|
|
async with Client(proxy_server) as client:
|
|
proxy_result = await client.read_resource("data://user/1")
|
|
assert proxy_result == result
|
|
|
|
|
|
class TestPrompts:
|
|
async def test_get_prompts_server_method(self, proxy_server: FastMCPProxy):
|
|
prompts = await proxy_server.get_prompts()
|
|
assert [p.name for p in prompts.values()] == Contains("welcome")
|
|
|
|
async def test_list_prompts_same_as_original(self, fastmcp_server, proxy_server):
|
|
async with Client(fastmcp_server) as client:
|
|
result = await client.list_prompts()
|
|
async with Client(proxy_server) as client:
|
|
proxy_result = await client.list_prompts()
|
|
assert proxy_result == result
|
|
|
|
async def test_render_prompt_same_as_original(
|
|
self, fastmcp_server: FastMCP, proxy_server: FastMCPProxy
|
|
):
|
|
async with Client(fastmcp_server) as client:
|
|
result = await client.get_prompt("welcome", {"name": "Alice"})
|
|
async with Client(proxy_server) as client:
|
|
proxy_result = await client.get_prompt("welcome", {"name": "Alice"})
|
|
assert proxy_result == result
|
|
|
|
async def test_render_prompt_calls_prompt(self, proxy_server):
|
|
async with Client(proxy_server) as client:
|
|
result = await client.get_prompt("welcome", {"name": "Alice"})
|
|
assert result.messages[0].role == "user"
|
|
assert result.messages[0].content.text == "Welcome to FastMCP, Alice!" # type: ignore[attr-defined]
|
|
|
|
|
|
async def test_proxy_handles_multiple_concurrent_tasks_correctly(
|
|
proxy_server: FastMCPProxy,
|
|
):
|
|
results = {}
|
|
|
|
async def get_and_store(name, coro):
|
|
results[name] = await coro()
|
|
|
|
async with create_task_group() as tg:
|
|
tg.start_soon(get_and_store, "prompts", proxy_server.get_prompts)
|
|
tg.start_soon(get_and_store, "resources", proxy_server.get_resources)
|
|
tg.start_soon(get_and_store, "tools", proxy_server.get_tools)
|
|
|
|
assert list(results) == Contains("resources", "prompts", "tools")
|
|
assert list(results["prompts"]) == Contains("welcome")
|
|
assert [r.uri for r in results["resources"].values()] == Contains(
|
|
AnyUrl("data://users"),
|
|
AnyUrl("resource://wave"),
|
|
)
|
|
assert [r.name for r in results["resources"].values()] == Contains(
|
|
"get_users", "wave"
|
|
)
|
|
assert list(results["tools"]) == Contains(
|
|
"greet", "add", "error_tool", "tool_without_description"
|
|
)
|