Spaces:
Running
Running
| import logging | |
| 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] = [] | |
| self.logger = logging.getLogger(__name__) | |
| # Backwards-compatible way to get the log level mapping | |
| if hasattr(logging, "getLevelNamesMapping"): | |
| # For Python 3.11+ | |
| self.LOGGING_LEVEL_MAP = logging.getLevelNamesMapping() # pyright: ignore [reportAttributeAccessIssue] | |
| else: | |
| # For older Python versions | |
| self.LOGGING_LEVEL_MAP = logging._nameToLevel | |
| async def handle_log(self, message: LogMessage) -> None: | |
| self.logs.append(message) | |
| level = self.LOGGING_LEVEL_MAP[message.level.upper()] | |
| msg = message.data.get("msg") | |
| extra = message.data.get("extra") | |
| self.logger.log(level, msg, extra=extra) | |
| def fastmcp_server(): | |
| mcp = FastMCP() | |
| async def log(context: Context) -> None: | |
| await context.info(message="hello?") | |
| 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, caplog): | |
| caplog.set_level(logging.INFO, logger=__name__) | |
| 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["msg"] == "hello?" | |
| assert log_handler.logs[0].level == "info" | |
| assert len(caplog.records) == 1 | |
| assert caplog.records[0].msg == "hello?" | |
| assert caplog.records[0].levelname == "INFO" | |
| async def test_echo_log(self, fastmcp_server: FastMCP, caplog): | |
| caplog.set_level(logging.INFO, logger=__name__) | |
| 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 | |
| assert len(caplog.records) == 1 | |
| await client.call_tool( | |
| "echo_log", {"message": "this is a warning log", "level": "warning"} | |
| ) | |
| assert len(log_handler.logs) == 2 | |
| assert len(caplog.records) == 2 | |
| assert log_handler.logs[0].data["msg"] == "this is a log" | |
| assert log_handler.logs[0].level == "info" | |
| assert log_handler.logs[1].data["msg"] == "this is a warning log" | |
| assert log_handler.logs[1].level == "warning" | |
| assert caplog.records[0].msg == "this is a log" | |
| assert caplog.records[0].levelname == "INFO" | |
| assert caplog.records[1].msg == "this is a warning log" | |
| assert caplog.records[1].levelname == "WARNING" | |
| class TestDefaultLogHandler: | |
| """Tests for default_log_handler bug fix (issue #1394).""" | |
| async def test_default_handler_routes_to_correct_levels(self): | |
| """Test that default_log_handler routes server logs to appropriate Python log levels.""" | |
| from unittest.mock import MagicMock, patch | |
| from mcp.types import LoggingMessageNotificationParams | |
| from fastmcp.client.logging import default_log_handler | |
| with patch("fastmcp.client.logging.logger") as mock_logger: | |
| # Set up mock methods | |
| mock_logger.debug = MagicMock() | |
| mock_logger.info = MagicMock() | |
| mock_logger.warning = MagicMock() | |
| mock_logger.error = MagicMock() | |
| mock_logger.critical = MagicMock() | |
| # Test each log level | |
| test_cases = [ | |
| ("debug", mock_logger.debug, "Debug message"), | |
| ("info", mock_logger.info, "Info message"), | |
| ("notice", mock_logger.info, "Notice message"), # notice -> info | |
| ("warning", mock_logger.warning, "Warning message"), | |
| ("error", mock_logger.error, "Error message"), | |
| ("critical", mock_logger.critical, "Critical message"), | |
| ("alert", mock_logger.critical, "Alert message"), # alert -> critical | |
| ( | |
| "emergency", | |
| mock_logger.critical, | |
| "Emergency message", | |
| ), # emergency -> critical | |
| ] | |
| for level, expected_method, msg in test_cases: | |
| # Reset mocks | |
| mock_logger.reset_mock() | |
| # Create log message | |
| log_msg = LoggingMessageNotificationParams( | |
| level=level, # type: ignore[arg-type] | |
| logger="test.logger", | |
| data={"msg": msg, "extra": {"test_key": "test_value"}}, | |
| ) | |
| # Call handler | |
| await default_log_handler(log_msg) | |
| # Verify correct method was called | |
| expected_method.assert_called_once_with( | |
| f"Server log: [test.logger] {msg}", extra={"test_key": "test_value"} | |
| ) | |
| async def test_default_handler_without_logger_name(self): | |
| """Test that default_log_handler works when logger name is None.""" | |
| from unittest.mock import MagicMock, patch | |
| from mcp.types import LoggingMessageNotificationParams | |
| from fastmcp.client.logging import default_log_handler | |
| with patch("fastmcp.client.logging.logger") as mock_logger: | |
| mock_logger.info = MagicMock() | |
| log_msg = LoggingMessageNotificationParams( | |
| level="info", | |
| logger=None, | |
| data={"msg": "Message without logger", "extra": {}}, | |
| ) | |
| await default_log_handler(log_msg) | |
| mock_logger.info.assert_called_once_with( | |
| "Server log: Message without logger", extra={} | |
| ) | |
| async def test_default_handler_with_missing_msg(self): | |
| """Test that default_log_handler handles missing 'msg' gracefully.""" | |
| from unittest.mock import MagicMock, patch | |
| from mcp.types import LoggingMessageNotificationParams | |
| from fastmcp.client.logging import default_log_handler | |
| with patch("fastmcp.client.logging.logger") as mock_logger: | |
| mock_logger.info = MagicMock() | |
| log_msg = LoggingMessageNotificationParams( | |
| level="info", | |
| logger="test.logger", | |
| data={"extra": {"key": "value"}}, # Missing 'msg' key | |
| ) | |
| await default_log_handler(log_msg) | |
| # Should use str(message) as fallback | |
| mock_logger.info.assert_called_once() | |
| call_args = mock_logger.info.call_args | |
| assert "Server log:" in call_args[0][0] | |
| assert call_args[1]["extra"] == {"key": "value"} | |