fix: add tests for telemetry (#31)

Author: hsheth2

src/mcp_server_datahub/__main__.py

@@ -1,68 +1,17 @@
 import logging
-from typing import Any
 
 import click
-import mcp.types as mt
 from datahub.ingestion.graph.config import ClientMode
 from datahub.sdk.main_client import DataHubClient
 from datahub.telemetry import telemetry
-from datahub.utilities.perf_timer import PerfTimer
-from fastmcp.server.middleware import CallNext, Middleware, MiddlewareContext
 from fastmcp.server.middleware.logging import LoggingMiddleware
 from typing_extensions import Literal
 
+from mcp_server_datahub._telemetry import TelemetryMiddleware
 from mcp_server_datahub._version import __version__
 from mcp_server_datahub.mcp_server import mcp, with_datahub_client
 
 logging.basicConfig(level=logging.INFO)
-telemetry.telemetry_instance.add_global_property(
-    "mcp_server_datahub_version", __version__
-)
-
-logger = logging.getLogger(__name__)
-
-
-class TelemetryMiddleware(Middleware):
-    """Middleware that logs tool calls."""
-
-    async def on_call_tool(
-        self,
-        context: MiddlewareContext[mt.CallToolRequestParams],
-        call_next: CallNext[mt.CallToolRequestParams, mt.CallToolResult],
-    ) -> mt.CallToolResult:
-        telemetry_data: dict[str, Any] = {}
-        with PerfTimer() as timer:
-            telemetry_data = {
-                "tool": context.message.name,
-                "source": context.source,
-                "type": context.type,
-                "method": context.method,
-            }
-            try:
-                result = await call_next(context)
-
-                # BUG: The FastMCP type annotations seem to be incorrect.
-                # This method typically returns fastmcp.tools.tool.ToolResult.
-                if isinstance(result, mt.CallToolResult):
-                    telemetry_data["tool_result_is_error"] = result.isError
-                telemetry_data["tool_result_length"] = (
-                    sum(
-                        len(block.text)
-                        for block in result.content
-                        if isinstance(block, mt.TextContent)
-                    ),
-                )
-
-                return result
-
-            except Exception as e:
-                telemetry_data["tool_call_error"] = e.__class__.__name__
-                raise
-            finally:
-                telemetry_data["duration_seconds"] = timer.elapsed_seconds()
-                telemetry.telemetry_instance.ping(
-                    "mcp-server-tool-call", telemetry_data
-                )
 
 
 @click.command()

src/mcp_server_datahub/_telemetry.py

@@ -0,0 +1,55 @@
+from typing import Any
+
+import mcp.types as mt
+from datahub.telemetry import telemetry
+from datahub.utilities.perf_timer import PerfTimer
+from fastmcp.server.middleware import CallNext, Middleware, MiddlewareContext
+
+from mcp_server_datahub._version import __version__
+
+telemetry.telemetry_instance.add_global_property(
+    "mcp_server_datahub_version", __version__
+)
+
+
+class TelemetryMiddleware(Middleware):
+    """Middleware that logs tool calls."""
+
+    async def on_call_tool(
+        self,
+        context: MiddlewareContext[mt.CallToolRequestParams],
+        call_next: CallNext[mt.CallToolRequestParams, mt.CallToolResult],
+    ) -> mt.CallToolResult:
+        telemetry_data: dict[str, Any] = {}
+        with PerfTimer() as timer:
+            telemetry_data = {
+                "tool": context.message.name,
+                "source": context.source,
+                "type": context.type,
+                "method": context.method,
+            }
+            try:
+                result = await call_next(context)
+
+                # BUG: The FastMCP type annotations seem to be incorrect.
+                # This method typically returns fastmcp.tools.tool.ToolResult.
+                if isinstance(result, mt.CallToolResult):
+                    telemetry_data["tool_result_is_error"] = result.isError
+                telemetry_data["tool_result_length"] = (
+                    sum(
+                        len(block.text)
+                        for block in result.content
+                        if isinstance(block, mt.TextContent)
+                    ),
+                )
+
+                return result
+
+            except Exception as e:
+                telemetry_data["tool_call_error"] = e.__class__.__name__
+                raise
+            finally:
+                telemetry_data["duration_seconds"] = timer.elapsed_seconds()
+                telemetry.telemetry_instance.ping(
+                    "mcp-server-tool-call", telemetry_data
+                )

tests/conftest.py

@@ -1,5 +1,9 @@
+import os
+
 import pytest
 
+os.environ["DATAHUB_TELEMETRY_ENABLED"] = "false"
+
 
 @pytest.fixture
 def anyio_backend() -> str:

tests/test_mcp_server.py

@@ -5,6 +5,7 @@
 from datahub.sdk.main_client import DataHubClient
 from fastmcp import Client
 
+from mcp_server_datahub._telemetry import TelemetryMiddleware
 from mcp_server_datahub.mcp_server import (
     get_dataset_queries,
     get_entity,
@@ -16,6 +17,10 @@
 _test_urn = "urn:li:dataset:(urn:li:dataPlatform:snowflake,long_tail_companions.analytics.pet_details,PROD)"
 _test_domain = "urn:li:domain:0da1ef03-8870-45db-9f47-ef4f592f095c"
 
+# Add telemetry middleware to the MCP server.
+# This way our tests also validate that the telemetry generation does not break anything else.
+mcp.add_middleware(TelemetryMiddleware())
+
 
 @pytest.fixture(autouse=True, scope="session")
 def setup_client() -> Iterable[None]: