mirror of
https://github.com/PrefectHQ/fastmcp.git
synced 2026-08-13 09:09:10 +02:00
516 lines
20 KiB
Python
516 lines
20 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, StreamableHttpTransport
|
|
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 isinstance(proxy.client.transport, StreamableHttpTransport)
|
|
assert proxy.client.transport.url == "http://example.com/mcp/"
|
|
|
|
|
|
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="This is a test error"):
|
|
async with Client(proxy_server) as client:
|
|
await client.call_tool("error_tool", {})
|
|
|
|
async def test_proxy_can_overwrite_proxied_tool(self, proxy_server):
|
|
"""
|
|
Test that a tool defined on the proxy can overwrite the proxied tool with the same name.
|
|
"""
|
|
|
|
@proxy_server.tool
|
|
def greet(name: str, extra: str = "extra") -> str:
|
|
return f"Overwritten, {name}! {extra}"
|
|
|
|
async with Client(proxy_server) as client:
|
|
result = await client.call_tool("greet", {"name": "Marvin", "extra": "abc"})
|
|
assert result[0].text == "Overwritten, Marvin! abc" # type: ignore[attr-defined]
|
|
|
|
async def test_proxy_errors_if_overwritten_tool_is_disabled(self, proxy_server):
|
|
"""
|
|
Test that a tool defined on the proxy is not listed if it is disabled,
|
|
and it doesn't fall back to the proxied tool with the same name
|
|
"""
|
|
|
|
@proxy_server.tool(enabled=False)
|
|
def greet(name: str, extra: str = "extra") -> str:
|
|
return f"Overwritten, {name}! {extra}"
|
|
|
|
async with Client(proxy_server) as client:
|
|
with pytest.raises(ToolError, match="Unknown tool"):
|
|
await client.call_tool("greet", {"name": "Marvin", "extra": "abc"})
|
|
|
|
async def test_proxy_can_list_overwritten_tool(self, proxy_server):
|
|
"""
|
|
Test that a tool defined on the proxy is listed instead of the proxied tool
|
|
"""
|
|
|
|
@proxy_server.tool
|
|
def greet(name: str, extra: str = "extra") -> str:
|
|
return f"Overwritten, {name}! {extra}"
|
|
|
|
async with Client(proxy_server) as client:
|
|
tools = await client.list_tools()
|
|
greet_tool = next(t for t in tools if t.name == "greet")
|
|
assert "extra" in greet_tool.inputSchema["properties"]
|
|
|
|
async def test_proxy_can_list_overwritten_tool_if_disabled(self, proxy_server):
|
|
"""
|
|
Test that a tool defined on the proxy is not listed if it is disabled,
|
|
and it doesn't fall back to the proxied tool with the same name
|
|
"""
|
|
|
|
@proxy_server.tool(enabled=False)
|
|
def greet(name: str, extra: str = "extra") -> str:
|
|
return f"Overwritten, {name}! {extra}"
|
|
|
|
async with Client(proxy_server) as client:
|
|
tools = await client.list_tools()
|
|
assert not any(t.name == "greet" for t in tools)
|
|
|
|
|
|
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")
|
|
|
|
async def test_proxy_can_overwrite_proxied_resource(self, proxy_server):
|
|
"""
|
|
Test that a resource defined on the proxy can overwrite the proxied resource with the same URI.
|
|
"""
|
|
|
|
@proxy_server.resource(uri="resource://wave")
|
|
def overwritten_wave() -> str:
|
|
return "Overwritten wave! 🌊"
|
|
|
|
async with Client(proxy_server) as client:
|
|
result = await client.read_resource("resource://wave")
|
|
assert result[0].text == "Overwritten wave! 🌊" # type: ignore[attr-defined]
|
|
|
|
async def test_proxy_errors_if_overwritten_resource_is_disabled(self, proxy_server):
|
|
"""
|
|
Test that a resource defined on the proxy is not accessible if it is disabled,
|
|
and it doesn't fall back to the proxied resource with the same URI
|
|
"""
|
|
|
|
@proxy_server.resource(uri="resource://wave", enabled=False)
|
|
def overwritten_wave() -> str:
|
|
return "Overwritten wave! 🌊"
|
|
|
|
async with Client(proxy_server) as client:
|
|
with pytest.raises(McpError, match="Unknown resource"):
|
|
await client.read_resource("resource://wave")
|
|
|
|
async def test_proxy_can_list_overwritten_resource(self, proxy_server):
|
|
"""
|
|
Test that a resource defined on the proxy is listed instead of the proxied resource
|
|
"""
|
|
|
|
@proxy_server.resource(uri="resource://wave", name="overwritten_wave")
|
|
def overwritten_wave() -> str:
|
|
return "Overwritten wave! 🌊"
|
|
|
|
async with Client(proxy_server) as client:
|
|
resources = await client.list_resources()
|
|
wave_resource = next(
|
|
r for r in resources if str(r.uri) == "resource://wave"
|
|
)
|
|
assert wave_resource.name == "overwritten_wave"
|
|
|
|
async def test_proxy_can_list_overwritten_resource_if_disabled(self, proxy_server):
|
|
"""
|
|
Test that a resource defined on the proxy is not listed if it is disabled,
|
|
and it doesn't fall back to the proxied resource with the same URI
|
|
"""
|
|
|
|
@proxy_server.resource(uri="resource://wave", enabled=False)
|
|
def overwritten_wave() -> str:
|
|
return "Overwritten wave! 🌊"
|
|
|
|
async with Client(proxy_server) as client:
|
|
resources = await client.list_resources()
|
|
wave_resources = [r for r in resources if str(r.uri) == "resource://wave"]
|
|
assert len(wave_resources) == 0
|
|
|
|
|
|
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
|
|
|
|
async def test_proxy_can_overwrite_proxied_resource_template(self, proxy_server):
|
|
"""
|
|
Test that a resource template defined on the proxy can overwrite the proxied template with the same URI template.
|
|
"""
|
|
|
|
@proxy_server.resource(uri="data://user/{user_id}", name="overwritten_get_user")
|
|
def overwritten_get_user(user_id: str) -> dict[str, Any]:
|
|
return {
|
|
"id": user_id,
|
|
"name": "Overwritten User",
|
|
"active": True,
|
|
"extra": "data",
|
|
}
|
|
|
|
async with Client(proxy_server) as client:
|
|
result = await client.read_resource("data://user/1")
|
|
user_data = json.loads(result[0].text) # type: ignore[attr-defined]
|
|
assert user_data["name"] == "Overwritten User"
|
|
assert user_data["extra"] == "data"
|
|
|
|
async def test_proxy_errors_if_overwritten_resource_template_is_disabled(
|
|
self, proxy_server
|
|
):
|
|
"""
|
|
Test that a resource template defined on the proxy is not accessible if it is disabled,
|
|
and it doesn't fall back to the proxied template with the same URI template
|
|
"""
|
|
|
|
@proxy_server.resource(uri="data://user/{user_id}", enabled=False)
|
|
def overwritten_get_user(user_id: str) -> dict[str, Any]:
|
|
return {"id": user_id, "name": "Overwritten User", "active": True}
|
|
|
|
async with Client(proxy_server) as client:
|
|
with pytest.raises(McpError, match="Unknown resource"):
|
|
await client.read_resource("data://user/1")
|
|
|
|
async def test_proxy_can_list_overwritten_resource_template(self, proxy_server):
|
|
"""
|
|
Test that a resource template defined on the proxy is listed instead of the proxied template
|
|
"""
|
|
|
|
@proxy_server.resource(uri="data://user/{user_id}", name="overwritten_get_user")
|
|
def overwritten_get_user(user_id: str) -> dict[str, Any]:
|
|
return {"id": user_id, "name": "Overwritten User", "active": True}
|
|
|
|
async with Client(proxy_server) as client:
|
|
templates = await client.list_resource_templates()
|
|
user_template = next(
|
|
t for t in templates if t.uriTemplate == "data://user/{user_id}"
|
|
)
|
|
assert user_template.name == "overwritten_get_user"
|
|
|
|
async def test_proxy_can_list_overwritten_resource_template_if_disabled(
|
|
self, proxy_server
|
|
):
|
|
"""
|
|
Test that a resource template defined on the proxy is not listed if it is disabled,
|
|
and it doesn't fall back to the proxied template with the same URI template
|
|
"""
|
|
|
|
@proxy_server.resource(uri="data://user/{user_id}", enabled=False)
|
|
def overwritten_get_user(user_id: str) -> dict[str, Any]:
|
|
return {"id": user_id, "name": "Overwritten User", "active": True}
|
|
|
|
async with Client(proxy_server) as client:
|
|
templates = await client.list_resource_templates()
|
|
user_templates = [
|
|
t for t in templates if t.uriTemplate == "data://user/{user_id}"
|
|
]
|
|
assert len(user_templates) == 0
|
|
|
|
|
|
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_can_overwrite_proxied_prompt(self, proxy_server):
|
|
"""
|
|
Test that a prompt defined on the proxy can overwrite the proxied prompt with the same name.
|
|
"""
|
|
|
|
@proxy_server.prompt
|
|
def welcome(name: str, extra: str = "friend") -> str:
|
|
return f"Overwritten welcome, {name}! You are my {extra}."
|
|
|
|
async with Client(proxy_server) as client:
|
|
result = await client.get_prompt(
|
|
"welcome", {"name": "Alice", "extra": "colleague"}
|
|
)
|
|
assert result.messages[0].role == "user"
|
|
assert (
|
|
result.messages[0].content.text # type: ignore[attr-defined]
|
|
== "Overwritten welcome, Alice! You are my colleague."
|
|
)
|
|
|
|
async def test_proxy_errors_if_overwritten_prompt_is_disabled(self, proxy_server):
|
|
"""
|
|
Test that a prompt defined on the proxy is not accessible if it is disabled,
|
|
and it doesn't fall back to the proxied prompt with the same name
|
|
"""
|
|
|
|
@proxy_server.prompt(enabled=False)
|
|
def welcome(name: str, extra: str = "friend") -> str:
|
|
return f"Overwritten welcome, {name}! You are my {extra}."
|
|
|
|
async with Client(proxy_server) as client:
|
|
with pytest.raises(McpError, match="Unknown prompt"):
|
|
await client.get_prompt("welcome", {"name": "Alice"})
|
|
|
|
async def test_proxy_can_list_overwritten_prompt(self, proxy_server):
|
|
"""
|
|
Test that a prompt defined on the proxy is listed instead of the proxied prompt
|
|
"""
|
|
|
|
@proxy_server.prompt
|
|
def welcome(name: str, extra: str = "friend") -> str:
|
|
return f"Overwritten welcome, {name}! You are my {extra}."
|
|
|
|
async with Client(proxy_server) as client:
|
|
prompts = await client.list_prompts()
|
|
welcome_prompt = next(p for p in prompts if p.name == "welcome")
|
|
# Check that the overwritten prompt has the additional 'extra' parameter
|
|
param_names = [arg.name for arg in welcome_prompt.arguments or []]
|
|
assert "extra" in param_names
|
|
|
|
async def test_proxy_can_list_overwritten_prompt_if_disabled(self, proxy_server):
|
|
"""
|
|
Test that a prompt defined on the proxy is not listed if it is disabled,
|
|
and it doesn't fall back to the proxied prompt with the same name
|
|
"""
|
|
|
|
@proxy_server.prompt(enabled=False)
|
|
def welcome(name: str, extra: str = "friend") -> str:
|
|
return f"Overwritten welcome, {name}! You are my {extra}."
|
|
|
|
async with Client(proxy_server) as client:
|
|
prompts = await client.list_prompts()
|
|
welcome_prompts = [p for p in prompts if p.name == "welcome"]
|
|
assert len(welcome_prompts) == 0
|
|
|
|
|
|
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"
|
|
)
|