mirror of
https://github.com/PrefectHQ/fastmcp.git
synced 2026-08-26 23:44:17 +02:00
286 lines
8.1 KiB
Python
286 lines
8.1 KiB
Python
import inspect
|
|
from pathlib import Path
|
|
|
|
from fastmcp.client.auth.bearer import BearerAuth
|
|
from fastmcp.client.auth.oauth import OAuthClientProvider
|
|
from fastmcp.client.client import Client
|
|
from fastmcp.client.logging import LogMessage
|
|
from fastmcp.client.transports import (
|
|
SSETransport,
|
|
StdioTransport,
|
|
StreamableHttpTransport,
|
|
)
|
|
from fastmcp.mcp_config import MCPConfig, RemoteMCPServer, StdioMCPServer
|
|
|
|
|
|
def test_parse_single_stdio_config():
|
|
config = {
|
|
"mcpServers": {
|
|
"test_server": {
|
|
"command": "echo",
|
|
"args": ["hello"],
|
|
}
|
|
}
|
|
}
|
|
mcp_config = MCPConfig.from_dict(config)
|
|
transport = mcp_config.mcpServers["test_server"].to_transport()
|
|
assert isinstance(transport, StdioTransport)
|
|
assert transport.command == "echo"
|
|
assert transport.args == ["hello"]
|
|
|
|
|
|
def test_parse_single_remote_config():
|
|
config = {
|
|
"mcpServers": {
|
|
"test_server": {
|
|
"url": "http://localhost:8000",
|
|
}
|
|
}
|
|
}
|
|
mcp_config = MCPConfig.from_dict(config)
|
|
transport = mcp_config.mcpServers["test_server"].to_transport()
|
|
assert isinstance(transport, StreamableHttpTransport)
|
|
assert transport.url == "http://localhost:8000"
|
|
|
|
|
|
def test_parse_remote_config_with_transport():
|
|
config = {
|
|
"mcpServers": {
|
|
"test_server": {
|
|
"url": "http://localhost:8000",
|
|
"transport": "sse",
|
|
}
|
|
}
|
|
}
|
|
mcp_config = MCPConfig.from_dict(config)
|
|
transport = mcp_config.mcpServers["test_server"].to_transport()
|
|
assert isinstance(transport, SSETransport)
|
|
assert transport.url == "http://localhost:8000"
|
|
|
|
|
|
def test_parse_remote_config_with_url_inference():
|
|
config = {
|
|
"mcpServers": {
|
|
"test_server": {
|
|
"url": "http://localhost:8000/sse/",
|
|
}
|
|
}
|
|
}
|
|
mcp_config = MCPConfig.from_dict(config)
|
|
transport = mcp_config.mcpServers["test_server"].to_transport()
|
|
assert isinstance(transport, SSETransport)
|
|
assert transport.url == "http://localhost:8000/sse/"
|
|
|
|
|
|
def test_parse_multiple_servers():
|
|
config = {
|
|
"mcpServers": {
|
|
"test_server": {
|
|
"url": "http://localhost:8000/sse/",
|
|
},
|
|
"test_server_2": {
|
|
"command": "echo",
|
|
"args": ["hello"],
|
|
"env": {"TEST": "test"},
|
|
},
|
|
}
|
|
}
|
|
mcp_config = MCPConfig.from_dict(config)
|
|
assert len(mcp_config.mcpServers) == 2
|
|
assert isinstance(mcp_config.mcpServers["test_server"], RemoteMCPServer)
|
|
assert isinstance(mcp_config.mcpServers["test_server"].to_transport(), SSETransport)
|
|
|
|
assert isinstance(mcp_config.mcpServers["test_server_2"], StdioMCPServer)
|
|
assert isinstance(
|
|
mcp_config.mcpServers["test_server_2"].to_transport(), StdioTransport
|
|
)
|
|
assert mcp_config.mcpServers["test_server_2"].command == "echo"
|
|
assert mcp_config.mcpServers["test_server_2"].args == ["hello"]
|
|
assert mcp_config.mcpServers["test_server_2"].env == {"TEST": "test"}
|
|
|
|
|
|
async def test_multi_client(tmp_path: Path):
|
|
server_script = inspect.cleandoc("""
|
|
from fastmcp import FastMCP
|
|
|
|
mcp = FastMCP()
|
|
|
|
@mcp.tool
|
|
def add(a: int, b: int) -> int:
|
|
return a + b
|
|
|
|
if __name__ == '__main__':
|
|
mcp.run()
|
|
""")
|
|
|
|
script_path = tmp_path / "test.py"
|
|
script_path.write_text(server_script)
|
|
|
|
config = {
|
|
"mcpServers": {
|
|
"test_1": {
|
|
"command": "python",
|
|
"args": [str(script_path)],
|
|
},
|
|
"test_2": {
|
|
"command": "python",
|
|
"args": [str(script_path)],
|
|
},
|
|
}
|
|
}
|
|
|
|
client = Client(config)
|
|
|
|
async with client:
|
|
tools = await client.list_tools()
|
|
assert len(tools) == 2
|
|
|
|
result_1 = await client.call_tool("test_1_add", {"a": 1, "b": 2})
|
|
result_2 = await client.call_tool("test_2_add", {"a": 1, "b": 2})
|
|
assert result_1.data == 3
|
|
assert result_2.data == 3
|
|
|
|
|
|
async def test_remote_config_default_no_auth():
|
|
config = {
|
|
"mcpServers": {
|
|
"test_server": {
|
|
"url": "http://localhost:8000",
|
|
}
|
|
}
|
|
}
|
|
client = Client(config)
|
|
assert isinstance(client.transport.transport, StreamableHttpTransport)
|
|
assert client.transport.transport.auth is None
|
|
|
|
|
|
async def test_remote_config_with_auth_token():
|
|
config = {
|
|
"mcpServers": {
|
|
"test_server": {
|
|
"url": "http://localhost:8000",
|
|
"auth": "test_token",
|
|
}
|
|
}
|
|
}
|
|
client = Client(config)
|
|
assert isinstance(client.transport.transport, StreamableHttpTransport)
|
|
assert isinstance(client.transport.transport.auth, BearerAuth)
|
|
assert client.transport.transport.auth.token.get_secret_value() == "test_token"
|
|
|
|
|
|
async def test_remote_config_sse_with_auth_token():
|
|
config = {
|
|
"mcpServers": {
|
|
"test_server": {
|
|
"url": "http://localhost:8000/sse/",
|
|
"auth": "test_token",
|
|
}
|
|
}
|
|
}
|
|
client = Client(config)
|
|
assert isinstance(client.transport.transport, SSETransport)
|
|
assert isinstance(client.transport.transport.auth, BearerAuth)
|
|
assert client.transport.transport.auth.token.get_secret_value() == "test_token"
|
|
|
|
|
|
async def test_remote_config_with_oauth_literal():
|
|
config = {
|
|
"mcpServers": {
|
|
"test_server": {
|
|
"url": "http://localhost:8000",
|
|
"auth": "oauth",
|
|
}
|
|
}
|
|
}
|
|
client = Client(config)
|
|
assert isinstance(client.transport.transport, StreamableHttpTransport)
|
|
assert isinstance(client.transport.transport.auth, OAuthClientProvider)
|
|
|
|
|
|
async def test_multi_client_with_logging(tmp_path: Path):
|
|
"""
|
|
Tests that logging is properly forwarded to the ultimate client.
|
|
"""
|
|
server_script = inspect.cleandoc("""
|
|
from fastmcp import FastMCP, Context
|
|
|
|
mcp = FastMCP()
|
|
|
|
@mcp.tool
|
|
async def log_test(message: str, ctx: Context) -> int:
|
|
await ctx.log(message)
|
|
return 42
|
|
|
|
if __name__ == '__main__':
|
|
mcp.run()
|
|
""")
|
|
|
|
script_path = tmp_path / "test.py"
|
|
script_path.write_text(server_script)
|
|
|
|
config = {
|
|
"mcpServers": {
|
|
"test_server": {
|
|
"command": "python",
|
|
"args": [str(script_path)],
|
|
},
|
|
"test_server_2": {
|
|
"command": "python",
|
|
"args": [str(script_path)],
|
|
},
|
|
}
|
|
}
|
|
|
|
MESSAGES = []
|
|
|
|
async def log_handler(message: LogMessage):
|
|
MESSAGES.append(message)
|
|
|
|
async with Client(config, log_handler=log_handler) as client:
|
|
result = await client.call_tool("test_server_log_test", {"message": "test 42"})
|
|
assert result.data == 42
|
|
assert len(MESSAGES) == 1
|
|
assert MESSAGES[0].data == "test 42"
|
|
|
|
|
|
async def test_multi_client_with_elicitation(tmp_path: Path):
|
|
"""
|
|
Tests that elicitation is properly forwarded to the ultimate client.
|
|
"""
|
|
server_script = inspect.cleandoc("""
|
|
from fastmcp import FastMCP, Context
|
|
|
|
mcp = FastMCP()
|
|
|
|
@mcp.tool
|
|
async def elicit_test(ctx: Context) -> int:
|
|
result = await ctx.elicit('Pick a number', response_type=int)
|
|
return result.data
|
|
|
|
if __name__ == '__main__':
|
|
mcp.run()
|
|
""")
|
|
|
|
script_path = tmp_path / "test.py"
|
|
script_path.write_text(server_script)
|
|
|
|
config = {
|
|
"mcpServers": {
|
|
"test_server": {
|
|
"command": "python",
|
|
"args": [str(script_path)],
|
|
},
|
|
"test_server_2": {
|
|
"command": "python",
|
|
"args": [str(script_path)],
|
|
},
|
|
}
|
|
}
|
|
|
|
async def elicitation_handler(message, response_type, params, ctx):
|
|
return response_type(value=42)
|
|
|
|
async with Client(config, elicitation_handler=elicitation_handler) as client:
|
|
result = await client.call_tool("test_server_elicit_test", {})
|
|
assert result.data == 42
|