From 70598a494423ecc7ec794b723518d8a3d165d041 Mon Sep 17 00:00:00 2001 From: Alexsander Hamir Date: Sat, 24 Jan 2026 11:04:34 -0800 Subject: [PATCH] Add Datadog mock client support - Create datadog_mock_client.py following GCS, Langfuse, and LangSmith patterns - Add mock mode detection via DATADOG_MOCK environment variable - Intercept Datadog API calls via AsyncHTTPHandler.post and httpx.Client.post patching - Add verbose logging throughout mock implementation - Update DataDogLogger and DataDogLLMObsLogger to initialize mock client when mock mode enabled - Supports both async and sync logging paths - Supports configurable mock latency via DATADOG_MOCK_LATENCY_MS --- litellm/integrations/datadog/datadog.py | 28 ++- .../integrations/datadog/datadog_llm_obs.py | 25 ++- .../datadog/datadog_mock_client.py | 176 ++++++++++++++++++ 3 files changed, 221 insertions(+), 8 deletions(-) create mode 100644 litellm/integrations/datadog/datadog_mock_client.py diff --git a/litellm/integrations/datadog/datadog.py b/litellm/integrations/datadog/datadog.py index 503e8d8c87a..08029f0da0a 100644 --- a/litellm/integrations/datadog/datadog.py +++ b/litellm/integrations/datadog/datadog.py @@ -27,6 +27,10 @@ import litellm from litellm._logging import verbose_logger from litellm._uuid import uuid from litellm.integrations.custom_batch_logger import CustomBatchLogger +from litellm.integrations.datadog.datadog_mock_client import ( + should_use_datadog_mock, + create_mock_datadog_client, +) from litellm.integrations.datadog.datadog_handler import ( get_datadog_hostname, get_datadog_service, @@ -80,6 +84,12 @@ class DataDogLogger( """ try: verbose_logger.debug("Datadog: in init datadog logger") + + self.is_mock_mode = should_use_datadog_mock() + + if self.is_mock_mode: + create_mock_datadog_client() + verbose_logger.debug("[DATADOG MOCK] Datadog logger initialized in mock mode") ######################################################### # Handle datadog_params set as litellm.datadog_params @@ -229,6 +239,9 @@ class DataDogLogger( len(self.log_queue), self.intake_url, ) + + if self.is_mock_mode: + verbose_logger.debug("[DATADOG MOCK] Mock mode enabled - API calls will be intercepted") response = await self.async_send_compressed_data(self.log_queue) if response.status_code == 413: @@ -241,11 +254,16 @@ class DataDogLogger( f"Response from datadog API status_code: {response.status_code}, text: {response.text}" ) - verbose_logger.debug( - "Datadog: Response from datadog API status_code: %s, text: %s", - response.status_code, - response.text, - ) + if self.is_mock_mode: + verbose_logger.debug( + f"[DATADOG MOCK] Batch of {len(self.log_queue)} events successfully mocked" + ) + else: + verbose_logger.debug( + "Datadog: Response from datadog API status_code: %s, text: %s", + response.status_code, + response.text, + ) except Exception as e: verbose_logger.exception( f"Datadog Error sending batch API - {str(e)}\n{traceback.format_exc()}" diff --git a/litellm/integrations/datadog/datadog_llm_obs.py b/litellm/integrations/datadog/datadog_llm_obs.py index 6ffdbc0a005..9291dfc2857 100644 --- a/litellm/integrations/datadog/datadog_llm_obs.py +++ b/litellm/integrations/datadog/datadog_llm_obs.py @@ -18,6 +18,10 @@ import httpx import litellm from litellm._logging import verbose_logger from litellm.integrations.custom_batch_logger import CustomBatchLogger +from litellm.integrations.datadog.datadog_mock_client import ( + should_use_datadog_mock, + create_mock_datadog_client, +) from litellm.integrations.datadog.datadog_handler import ( get_datadog_service, get_datadog_tags, @@ -43,6 +47,13 @@ class DataDogLLMObsLogger(CustomBatchLogger): def __init__(self, **kwargs): try: verbose_logger.debug("DataDogLLMObs: Initializing logger") + + self.is_mock_mode = should_use_datadog_mock() + + if self.is_mock_mode: + create_mock_datadog_client() + verbose_logger.debug("[DATADOG MOCK] DataDogLLMObs logger initialized in mock mode") + if os.getenv("DD_API_KEY", None) is None: raise Exception("DD_API_KEY is not set, set 'DD_API_KEY=<>'") if os.getenv("DD_SITE", None) is None: @@ -139,6 +150,9 @@ class DataDogLLMObsLogger(CustomBatchLogger): verbose_logger.debug( f"DataDogLLMObs: Flushing {len(self.log_queue)} events" ) + + if self.is_mock_mode: + verbose_logger.debug("[DATADOG MOCK] Mock mode enabled - API calls will be intercepted") # Prepare the payload payload = { @@ -178,9 +192,14 @@ class DataDogLLMObsLogger(CustomBatchLogger): f"DataDogLLMObs: Unexpected response - status_code: {response.status_code}, text: {response.text}" ) - verbose_logger.debug( - f"DataDogLLMObs: Successfully sent batch - status_code: {response.status_code}" - ) + if self.is_mock_mode: + verbose_logger.debug( + f"[DATADOG MOCK] Batch of {len(self.log_queue)} events successfully mocked" + ) + else: + verbose_logger.debug( + f"DataDogLLMObs: Successfully sent batch - status_code: {response.status_code}" + ) self.log_queue.clear() except httpx.HTTPStatusError as e: verbose_logger.exception( diff --git a/litellm/integrations/datadog/datadog_mock_client.py b/litellm/integrations/datadog/datadog_mock_client.py new file mode 100644 index 00000000000..80c1003510d --- /dev/null +++ b/litellm/integrations/datadog/datadog_mock_client.py @@ -0,0 +1,176 @@ +""" +Mock client for Datadog integration testing. + +This module intercepts Datadog API calls and returns successful mock responses, +allowing full code execution without making actual network calls. + +Usage: + Set DATADOG_MOCK=true in environment variables or config to enable mock mode. +""" + +import httpx +import json +import asyncio +from datetime import timedelta +from typing import Dict, Optional + +from litellm._logging import verbose_logger + +# Store original methods for restoration +_original_async_handler_post = None +_original_sync_client_post = None + +# Track if mocks have been initialized to avoid duplicate initialization +_mocks_initialized = False + +# Default mock latency in seconds (simulates network round-trip) +# Typical Datadog API calls take 50-150ms +_MOCK_LATENCY_SECONDS = float(__import__("os").getenv("DATADOG_MOCK_LATENCY_MS", "100")) / 1000.0 + + +class MockDatadogResponse: + """Mock httpx.Response that satisfies Datadog API requirements.""" + + def __init__(self, status_code: int = 202, json_data: Optional[Dict] = None, url: Optional[str] = None, elapsed_seconds: float = 0.0): + self.status_code = status_code + self._json_data = json_data or {"status": "ok"} + self.headers = httpx.Headers({}) + self.is_success = status_code < 400 + self.is_error = status_code >= 400 + self.is_redirect = 300 <= status_code < 400 + self.url = httpx.URL(url) if url else httpx.URL("") + # Set realistic elapsed time based on mock latency + elapsed_time = elapsed_seconds if elapsed_seconds > 0 else _MOCK_LATENCY_SECONDS + self.elapsed = timedelta(seconds=elapsed_time) + self._text = json.dumps(self._json_data) if json_data else "" + self._content = self._text.encode("utf-8") + + @property + def text(self) -> str: + """Return response text.""" + return self._text + + @property + def content(self) -> bytes: + """Return response content.""" + return self._content + + def json(self) -> Dict: + """Return JSON response data.""" + return self._json_data + + def read(self) -> bytes: + """Read response content.""" + return self._content + + def raise_for_status(self): + """Raise exception for error status codes.""" + if self.status_code >= 400: + raise Exception(f"HTTP {self.status_code}") + + +def _is_datadog_url(url) -> bool: + """Check if URL is a Datadog domain.""" + try: + parsed_url = httpx.URL(url) if isinstance(url, str) else url + hostname = parsed_url.host or "" + + return ( + hostname.endswith(".datadoghq.com") or + hostname == "datadoghq.com" or + "datadoghq.com" in hostname or + (hostname in ("localhost", "127.0.0.1") and "datadog" in str(parsed_url).lower()) + ) + except Exception: + return False + + +async def _mock_async_handler_post(self, url, data=None, json=None, params=None, headers=None, timeout=None, stream=False, logging_obj=None, files=None, content=None): + """Monkey-patched AsyncHTTPHandler.post that intercepts Datadog calls.""" + # Only mock Datadog API calls + if isinstance(url, str) and _is_datadog_url(url): + verbose_logger.info(f"[DATADOG MOCK] POST to {url}") + # Simulate network latency + await asyncio.sleep(_MOCK_LATENCY_SECONDS) + return MockDatadogResponse( + status_code=202, + json_data={"status": "ok"}, + url=url, + elapsed_seconds=_MOCK_LATENCY_SECONDS + ) + # For non-Datadog calls, use original method + if _original_async_handler_post is not None: + return await _original_async_handler_post(self, url=url, data=data, json=json, params=params, headers=headers, timeout=timeout, stream=stream, logging_obj=logging_obj, files=files, content=content) + # Fallback: if original not set, raise error + raise RuntimeError("Original AsyncHTTPHandler.post not available") + + +def _mock_sync_client_post(self, url, **kwargs): + """Monkey-patched httpx.Client.post that intercepts Datadog calls.""" + if _is_datadog_url(url): + verbose_logger.info(f"[DATADOG MOCK] POST to {url} (sync)") + return MockDatadogResponse(status_code=202, json_data={"status": "ok"}, url=url, elapsed_seconds=_MOCK_LATENCY_SECONDS) + + if _original_sync_client_post is not None: + return _original_sync_client_post(self, url, **kwargs) + + +def create_mock_datadog_client(): + """ + Monkey-patch AsyncHTTPHandler.post and httpx.Client.post to intercept Datadog calls. + + AsyncHTTPHandler is used by LiteLLM's get_async_httpx_client() which is what + DataDogLogger and DataDogLLMObsLogger use for making API calls. + + httpx.Client is used for sync logging in DataDogLogger. + + This function is idempotent - it only initializes mocks once, even if called multiple times. + """ + global _original_async_handler_post, _original_sync_client_post + global _mocks_initialized + + # If already initialized, skip + if _mocks_initialized: + return + + verbose_logger.debug("[DATADOG MOCK] Initializing Datadog mock client...") + + # Patch AsyncHTTPHandler.post (used by LiteLLM's custom httpx handler) + if _original_async_handler_post is None: + from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler + _original_async_handler_post = AsyncHTTPHandler.post + AsyncHTTPHandler.post = _mock_async_handler_post # type: ignore + verbose_logger.debug("[DATADOG MOCK] Patched AsyncHTTPHandler.post") + + # Patch httpx.Client.post (used for sync logging) + if _original_sync_client_post is None: + _original_sync_client_post = httpx.Client.post + httpx.Client.post = _mock_sync_client_post # type: ignore + verbose_logger.debug("[DATADOG MOCK] Patched httpx.Client.post") + + verbose_logger.debug(f"[DATADOG MOCK] Mock latency set to {_MOCK_LATENCY_SECONDS*1000:.0f}ms") + verbose_logger.debug("[DATADOG MOCK] Datadog mock client initialization complete") + + _mocks_initialized = True + + +def should_use_datadog_mock() -> bool: + """ + Determine if Datadog should run in mock mode. + + Checks the DATADOG_MOCK environment variable. + + Returns: + bool: True if mock mode should be enabled + """ + import os + from litellm.secret_managers.main import str_to_bool + + mock_mode = os.getenv("DATADOG_MOCK", "false") + result = str_to_bool(mock_mode) + result = bool(result) if result is not None else False + + if result: + verbose_logger.info("Datadog Mock Mode: ENABLED - API calls will be mocked") + + return result