mirror of
https://github.com/PrefectHQ/fastmcp.git
synced 2026-08-15 10:09:11 +02:00
60 lines
1.8 KiB
Python
60 lines
1.8 KiB
Python
import pytest
|
|
from mcp import LoggingLevel
|
|
|
|
from fastmcp import Client, Context, FastMCP
|
|
from fastmcp.client.logging import LogMessage
|
|
|
|
|
|
class LogHandler:
|
|
def __init__(self):
|
|
self.logs: list[LogMessage] = []
|
|
|
|
async def handle_log(self, message: LogMessage) -> None:
|
|
self.logs.append(message)
|
|
|
|
|
|
@pytest.fixture
|
|
def fastmcp_server():
|
|
mcp = FastMCP()
|
|
|
|
@mcp.tool()
|
|
async def log(context: Context) -> None:
|
|
await context.info(message="hello?")
|
|
|
|
@mcp.tool()
|
|
async def echo_log(
|
|
message: str,
|
|
context: Context,
|
|
level: LoggingLevel | None = None,
|
|
logger: str | None = None,
|
|
) -> None:
|
|
await context.log(message=message, level=level)
|
|
|
|
return mcp
|
|
|
|
|
|
class TestClientLogs:
|
|
async def test_log(self, fastmcp_server: FastMCP):
|
|
log_handler = LogHandler()
|
|
async with Client(fastmcp_server, log_handler=log_handler.handle_log) as client:
|
|
await client.call_tool("log", {})
|
|
|
|
assert len(log_handler.logs) == 1
|
|
assert log_handler.logs[0].data == "hello?"
|
|
assert log_handler.logs[0].level == "info"
|
|
|
|
async def test_echo_log(self, fastmcp_server: FastMCP):
|
|
log_handler = LogHandler()
|
|
async with Client(fastmcp_server, log_handler=log_handler.handle_log) as client:
|
|
await client.call_tool("echo_log", {"message": "this is a log"})
|
|
|
|
assert len(log_handler.logs) == 1
|
|
await client.call_tool(
|
|
"echo_log", {"message": "this is a warning log", "level": "warning"}
|
|
)
|
|
assert len(log_handler.logs) == 2
|
|
|
|
assert log_handler.logs[0].data == "this is a log"
|
|
assert log_handler.logs[0].level == "info"
|
|
assert log_handler.logs[1].data == "this is a warning log"
|
|
assert log_handler.logs[1].level == "warning"
|