mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-22 00:31:44 +00:00
refactor: consolidate mock client logic into factory pattern
- Create mock_client_factory.py to centralize common mock HTTP client logic - Refactor GCS, Langfuse, LangSmith, and Datadog mock clients to use factory - Improve GET/DELETE mock accuracy for GCS (return valid StandardLoggingPayload) - Fix DELETE mock to return empty body (204 No Content) instead of JSON - Reduce code duplication across integration mock clients
This commit is contained in:
Alexsander Hamir
parent
70598a4944
commit
e5bea6dd28
5 changed files with 314 additions and 525 deletions
|
|
@ -8,169 +8,21 @@ 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.integrations.mock_client_factory import MockClientConfig, create_mock_client_factory
|
||||
|
||||
from litellm._logging import verbose_logger
|
||||
# Create mock client using factory
|
||||
_config = MockClientConfig(
|
||||
name="DATADOG",
|
||||
env_var="DATADOG_MOCK",
|
||||
default_latency_ms=100,
|
||||
default_status_code=202,
|
||||
default_json_data={"status": "ok"},
|
||||
url_matchers=[
|
||||
".datadoghq.com",
|
||||
"datadoghq.com",
|
||||
],
|
||||
patch_async_handler=True,
|
||||
patch_sync_client=True,
|
||||
)
|
||||
|
||||
# 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
|
||||
create_mock_datadog_client, should_use_datadog_mock = create_mock_client_factory(_config)
|
||||
|
|
|
|||
|
|
@ -15,13 +15,25 @@ from datetime import timedelta
|
|||
from typing import Dict, Optional
|
||||
|
||||
from litellm._logging import verbose_logger
|
||||
from litellm.integrations.mock_client_factory import MockClientConfig, create_mock_client_factory, MockResponse
|
||||
|
||||
# Store original methods for restoration
|
||||
_original_async_handler_post = None
|
||||
# Use factory for POST handler
|
||||
_config = MockClientConfig(
|
||||
name="GCS",
|
||||
env_var="GCS_MOCK",
|
||||
default_latency_ms=150,
|
||||
default_status_code=200,
|
||||
default_json_data={"kind": "storage#object", "name": "mock-object"},
|
||||
url_matchers=["storage.googleapis.com"],
|
||||
patch_async_handler=True,
|
||||
patch_sync_client=False,
|
||||
)
|
||||
|
||||
_create_mock_gcs_post, should_use_gcs_mock = create_mock_client_factory(_config)
|
||||
|
||||
# Store original methods for GET/DELETE (GCS-specific)
|
||||
_original_async_handler_get = None
|
||||
_original_async_handler_delete = None
|
||||
|
||||
# Track if mocks have been initialized to avoid duplicate initialization
|
||||
_mocks_initialized = False
|
||||
|
||||
# Default mock latency in seconds (simulates network round-trip)
|
||||
|
|
@ -29,84 +41,59 @@ _mocks_initialized = False
|
|||
_MOCK_LATENCY_SECONDS = float(__import__("os").getenv("GCS_MOCK_LATENCY_MS", "150")) / 1000.0
|
||||
|
||||
|
||||
class MockGCSResponse:
|
||||
"""Mock httpx.Response that satisfies GCS API requirements."""
|
||||
|
||||
def __init__(self, status_code: int = 200, 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 {"kind": "storage#object", "name": "mock-object"}
|
||||
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)
|
||||
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}")
|
||||
|
||||
|
||||
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 GCS calls."""
|
||||
# Only mock GCS API calls
|
||||
if isinstance(url, str) and "storage.googleapis.com" in url:
|
||||
verbose_logger.info(f"[GCS MOCK] POST to {url}")
|
||||
# Simulate network latency
|
||||
await asyncio.sleep(_MOCK_LATENCY_SECONDS)
|
||||
return MockGCSResponse(
|
||||
status_code=200,
|
||||
json_data={"kind": "storage#object", "name": "mock-object"},
|
||||
url=url,
|
||||
elapsed_seconds=_MOCK_LATENCY_SECONDS
|
||||
)
|
||||
# For non-GCS 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")
|
||||
|
||||
|
||||
async def _mock_async_handler_get(self, url, params=None, headers=None, follow_redirects=None):
|
||||
"""Monkey-patched AsyncHTTPHandler.get that intercepts GCS calls."""
|
||||
# Only mock GCS API calls
|
||||
if isinstance(url, str) and "storage.googleapis.com" in url:
|
||||
verbose_logger.info(f"[GCS MOCK] GET to {url}")
|
||||
# Simulate network latency
|
||||
await asyncio.sleep(_MOCK_LATENCY_SECONDS)
|
||||
return MockGCSResponse(
|
||||
status_code=200,
|
||||
json_data={"data": "mock-log-data"},
|
||||
# Return a minimal but valid StandardLoggingPayload JSON string as bytes
|
||||
# This matches what GCS returns when downloading with ?alt=media
|
||||
mock_payload = {
|
||||
"id": "mock-request-id",
|
||||
"trace_id": "mock-trace-id",
|
||||
"call_type": "completion",
|
||||
"stream": False,
|
||||
"response_cost": 0.0,
|
||||
"status": "success",
|
||||
"status_fields": {"llm_api_status": "success"},
|
||||
"custom_llm_provider": "mock",
|
||||
"total_tokens": 0,
|
||||
"prompt_tokens": 0,
|
||||
"completion_tokens": 0,
|
||||
"startTime": 0.0,
|
||||
"endTime": 0.0,
|
||||
"completionStartTime": 0.0,
|
||||
"response_time": 0.0,
|
||||
"model_map_information": {"model": "mock-model"},
|
||||
"model": "mock-model",
|
||||
"model_id": None,
|
||||
"model_group": None,
|
||||
"api_base": "https://api.mock.com",
|
||||
"metadata": {},
|
||||
"cache_hit": None,
|
||||
"cache_key": None,
|
||||
"saved_cache_cost": 0.0,
|
||||
"request_tags": [],
|
||||
"end_user": None,
|
||||
"requester_ip_address": None,
|
||||
"messages": None,
|
||||
"response": None,
|
||||
"error_str": None,
|
||||
"error_information": None,
|
||||
"model_parameters": {},
|
||||
"hidden_params": {},
|
||||
"guardrail_information": None,
|
||||
"standard_built_in_tools_params": None,
|
||||
}
|
||||
return MockResponse(
|
||||
status_code=200,
|
||||
json_data=mock_payload,
|
||||
url=url,
|
||||
elapsed_seconds=_MOCK_LATENCY_SECONDS
|
||||
)
|
||||
# For non-GCS calls, use original method
|
||||
if _original_async_handler_get is not None:
|
||||
return await _original_async_handler_get(self, url=url, params=params, headers=headers, follow_redirects=follow_redirects)
|
||||
# Fallback: if original not set, raise error
|
||||
raise RuntimeError("Original AsyncHTTPHandler.get not available")
|
||||
|
||||
|
||||
|
|
@ -115,18 +102,16 @@ async def _mock_async_handler_delete(self, url, data=None, json=None, params=Non
|
|||
# Only mock GCS API calls
|
||||
if isinstance(url, str) and "storage.googleapis.com" in url:
|
||||
verbose_logger.info(f"[GCS MOCK] DELETE to {url}")
|
||||
# Simulate network latency
|
||||
await asyncio.sleep(_MOCK_LATENCY_SECONDS)
|
||||
return MockGCSResponse(
|
||||
status_code=204,
|
||||
json_data={},
|
||||
# DELETE returns 204 No Content with empty body (not JSON)
|
||||
return MockResponse(
|
||||
status_code=204,
|
||||
json_data=None, # Empty body for DELETE
|
||||
url=url,
|
||||
elapsed_seconds=_MOCK_LATENCY_SECONDS
|
||||
)
|
||||
# For non-GCS calls, use original method
|
||||
if _original_async_handler_delete is not None:
|
||||
return await _original_async_handler_delete(self, url=url, data=data, json=json, params=params, headers=headers, timeout=timeout, stream=stream, content=content)
|
||||
# Fallback: if original not set, raise error
|
||||
raise RuntimeError("Original AsyncHTTPHandler.delete not available")
|
||||
|
||||
|
||||
|
|
@ -139,30 +124,26 @@ def create_mock_gcs_client():
|
|||
|
||||
This function is idempotent - it only initializes mocks once, even if called multiple times.
|
||||
"""
|
||||
global _original_async_handler_post, _original_async_handler_get, _original_async_handler_delete
|
||||
global _mocks_initialized
|
||||
global _original_async_handler_get, _original_async_handler_delete, _mocks_initialized
|
||||
|
||||
# If already initialized, skip
|
||||
# Use factory for POST handler
|
||||
_create_mock_gcs_post()
|
||||
|
||||
# If already initialized, skip GET/DELETE patching
|
||||
if _mocks_initialized:
|
||||
return
|
||||
|
||||
verbose_logger.debug("[GCS MOCK] Initializing GCS mock client...")
|
||||
verbose_logger.debug("[GCS MOCK] Initializing GCS GET/DELETE handlers...")
|
||||
|
||||
# Patch AsyncHTTPHandler methods (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("[GCS MOCK] Patched AsyncHTTPHandler.post")
|
||||
# Patch GET and DELETE handlers (GCS-specific)
|
||||
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler
|
||||
|
||||
if _original_async_handler_get is None:
|
||||
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler
|
||||
_original_async_handler_get = AsyncHTTPHandler.get
|
||||
AsyncHTTPHandler.get = _mock_async_handler_get # type: ignore
|
||||
verbose_logger.debug("[GCS MOCK] Patched AsyncHTTPHandler.get")
|
||||
|
||||
if _original_async_handler_delete is None:
|
||||
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler
|
||||
_original_async_handler_delete = AsyncHTTPHandler.delete
|
||||
AsyncHTTPHandler.delete = _mock_async_handler_delete # type: ignore
|
||||
verbose_logger.debug("[GCS MOCK] Patched AsyncHTTPHandler.delete")
|
||||
|
|
@ -212,25 +193,4 @@ def mock_vertex_auth_methods():
|
|||
verbose_logger.debug("[GCS MOCK] Patched Vertex AI auth methods")
|
||||
|
||||
|
||||
def should_use_gcs_mock() -> bool:
|
||||
"""
|
||||
Determine if GCS should run in mock mode.
|
||||
|
||||
Checks the GCS_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("GCS_MOCK", "false")
|
||||
result = str_to_bool(mock_mode)
|
||||
|
||||
# Ensure we return a bool, not None
|
||||
result = bool(result) if result is not None else False
|
||||
|
||||
if result:
|
||||
verbose_logger.info("GCS Mock Mode: ENABLED - API calls will be mocked")
|
||||
|
||||
return result
|
||||
# should_use_gcs_mock is already created by the factory
|
||||
|
|
|
|||
|
|
@ -9,113 +9,27 @@ Usage:
|
|||
"""
|
||||
|
||||
import httpx
|
||||
import json
|
||||
from datetime import timedelta
|
||||
from typing import Dict, Optional
|
||||
from litellm.integrations.mock_client_factory import MockClientConfig, create_mock_client_factory
|
||||
|
||||
from litellm._logging import verbose_logger
|
||||
|
||||
_original_httpx_post = None
|
||||
|
||||
# Default mock latency in seconds (simulates network round-trip)
|
||||
# Typical Langfuse API calls take 50-150ms
|
||||
_MOCK_LATENCY_SECONDS = float(__import__("os").getenv("LANGFUSE_MOCK_LATENCY_MS", "100")) / 1000.0
|
||||
|
||||
|
||||
class MockLangfuseResponse:
|
||||
"""Mock httpx.Response that satisfies Langfuse SDK requirements."""
|
||||
|
||||
def __init__(self, status_code: int = 200, 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": "success"}
|
||||
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)
|
||||
self._content = self._text.encode("utf-8")
|
||||
|
||||
@property
|
||||
def text(self) -> str:
|
||||
return self._text
|
||||
|
||||
@property
|
||||
def content(self) -> bytes:
|
||||
return self._content
|
||||
|
||||
def json(self) -> Dict:
|
||||
return self._json_data
|
||||
|
||||
def read(self) -> bytes:
|
||||
return self._content
|
||||
|
||||
def raise_for_status(self):
|
||||
if self.status_code >= 400:
|
||||
raise Exception(f"HTTP {self.status_code}")
|
||||
|
||||
|
||||
def _is_langfuse_url(url) -> bool:
|
||||
"""Check if URL is a Langfuse domain."""
|
||||
try:
|
||||
parsed_url = httpx.URL(url) if isinstance(url, str) else url
|
||||
hostname = parsed_url.host or ""
|
||||
|
||||
return (
|
||||
hostname.endswith(".langfuse.com") or
|
||||
hostname == "langfuse.com" or
|
||||
(hostname in ("localhost", "127.0.0.1") and "langfuse" in str(parsed_url).lower())
|
||||
)
|
||||
except Exception:
|
||||
return False
|
||||
|
||||
|
||||
def _mock_httpx_post(self, url, **kwargs):
|
||||
"""Monkey-patched httpx.Client.post that intercepts Langfuse calls."""
|
||||
if _is_langfuse_url(url):
|
||||
verbose_logger.info(f"[LANGFUSE MOCK] POST to {url}")
|
||||
return MockLangfuseResponse(status_code=200, json_data={"status": "success"}, url=url, elapsed_seconds=_MOCK_LATENCY_SECONDS)
|
||||
|
||||
if _original_httpx_post is not None:
|
||||
return _original_httpx_post(self, url, **kwargs)
|
||||
# Create mock client using factory
|
||||
_config = MockClientConfig(
|
||||
name="LANGFUSE",
|
||||
env_var="LANGFUSE_MOCK",
|
||||
default_latency_ms=100,
|
||||
default_status_code=200,
|
||||
default_json_data={"status": "success"},
|
||||
url_matchers=[
|
||||
".langfuse.com",
|
||||
"langfuse.com",
|
||||
],
|
||||
patch_async_handler=False,
|
||||
patch_sync_client=True,
|
||||
)
|
||||
|
||||
_create_mock_langfuse_client_internal, should_use_langfuse_mock = create_mock_client_factory(_config)
|
||||
|
||||
# Langfuse needs to return an httpx.Client instance
|
||||
def create_mock_langfuse_client():
|
||||
"""
|
||||
Monkey-patch httpx.Client.post to intercept Langfuse calls.
|
||||
|
||||
Returns a real httpx.Client instance - the monkey-patch intercepts all calls.
|
||||
"""
|
||||
global _original_httpx_post
|
||||
|
||||
if _original_httpx_post is None:
|
||||
_original_httpx_post = httpx.Client.post
|
||||
httpx.Client.post = _mock_httpx_post # type: ignore
|
||||
verbose_logger.debug("[LANGFUSE MOCK] Patched httpx.Client.post")
|
||||
|
||||
"""Create and return an httpx.Client instance - the monkey-patch intercepts all calls."""
|
||||
_create_mock_langfuse_client_internal()
|
||||
return httpx.Client()
|
||||
|
||||
|
||||
def should_use_langfuse_mock() -> bool:
|
||||
"""
|
||||
Determine if Langfuse should run in mock mode.
|
||||
|
||||
Checks the LANGFUSE_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("LANGFUSE_MOCK", "false")
|
||||
result = str_to_bool(mock_mode)
|
||||
result = bool(result) if result is not None else False
|
||||
|
||||
if result:
|
||||
verbose_logger.info("Langfuse Mock Mode: ENABLED - API calls will be mocked")
|
||||
|
||||
return result
|
||||
|
|
|
|||
|
|
@ -8,150 +8,22 @@ Usage:
|
|||
Set LANGSMITH_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.integrations.mock_client_factory import MockClientConfig, create_mock_client_factory
|
||||
|
||||
from litellm._logging import verbose_logger
|
||||
# Create mock client using factory
|
||||
_config = MockClientConfig(
|
||||
name="LANGSMITH",
|
||||
env_var="LANGSMITH_MOCK",
|
||||
default_latency_ms=100,
|
||||
default_status_code=200,
|
||||
default_json_data={"status": "success", "ids": ["mock-run-id"]},
|
||||
url_matchers=[
|
||||
".smith.langchain.com",
|
||||
"api.smith.langchain.com",
|
||||
"smith.langchain.com",
|
||||
],
|
||||
patch_async_handler=True,
|
||||
patch_sync_client=False,
|
||||
)
|
||||
|
||||
# Store original methods for restoration
|
||||
_original_async_handler_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 LangSmith API calls take 50-150ms
|
||||
_MOCK_LATENCY_SECONDS = float(__import__("os").getenv("LANGSMITH_MOCK_LATENCY_MS", "100")) / 1000.0
|
||||
|
||||
|
||||
class MockLangsmithResponse:
|
||||
"""Mock httpx.Response that satisfies LangSmith API requirements."""
|
||||
|
||||
def __init__(self, status_code: int = 200, 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": "success"}
|
||||
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)
|
||||
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_langsmith_url(url) -> bool:
|
||||
"""Check if URL is a LangSmith domain."""
|
||||
try:
|
||||
parsed_url = httpx.URL(url) if isinstance(url, str) else url
|
||||
hostname = parsed_url.host or ""
|
||||
|
||||
return (
|
||||
hostname.endswith(".smith.langchain.com") or
|
||||
hostname == "api.smith.langchain.com" or
|
||||
"smith.langchain.com" in hostname or
|
||||
(hostname in ("localhost", "127.0.0.1") and "langsmith" 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 LangSmith calls."""
|
||||
# Only mock LangSmith API calls
|
||||
if isinstance(url, str) and _is_langsmith_url(url):
|
||||
verbose_logger.info(f"[LANGSMITH MOCK] POST to {url}")
|
||||
# Simulate network latency
|
||||
await asyncio.sleep(_MOCK_LATENCY_SECONDS)
|
||||
return MockLangsmithResponse(
|
||||
status_code=200,
|
||||
json_data={"status": "success", "ids": ["mock-run-id"]},
|
||||
url=url,
|
||||
elapsed_seconds=_MOCK_LATENCY_SECONDS
|
||||
)
|
||||
# For non-LangSmith 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 create_mock_langsmith_client():
|
||||
"""
|
||||
Monkey-patch AsyncHTTPHandler.post to intercept LangSmith calls.
|
||||
|
||||
AsyncHTTPHandler is used by LiteLLM's get_async_httpx_client() which is what
|
||||
LangsmithLogger uses for making API calls.
|
||||
|
||||
This function is idempotent - it only initializes mocks once, even if called multiple times.
|
||||
"""
|
||||
global _original_async_handler_post
|
||||
global _mocks_initialized
|
||||
|
||||
# If already initialized, skip
|
||||
if _mocks_initialized:
|
||||
return
|
||||
|
||||
verbose_logger.debug("[LANGSMITH MOCK] Initializing LangSmith 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("[LANGSMITH MOCK] Patched AsyncHTTPHandler.post")
|
||||
|
||||
verbose_logger.debug(f"[LANGSMITH MOCK] Mock latency set to {_MOCK_LATENCY_SECONDS*1000:.0f}ms")
|
||||
verbose_logger.debug("[LANGSMITH MOCK] LangSmith mock client initialization complete")
|
||||
|
||||
_mocks_initialized = True
|
||||
|
||||
|
||||
def should_use_langsmith_mock() -> bool:
|
||||
"""
|
||||
Determine if LangSmith should run in mock mode.
|
||||
|
||||
Checks the LANGSMITH_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("LANGSMITH_MOCK", "false")
|
||||
result = str_to_bool(mock_mode)
|
||||
result = bool(result) if result is not None else False
|
||||
|
||||
if result:
|
||||
verbose_logger.info("LangSmith Mock Mode: ENABLED - API calls will be mocked")
|
||||
|
||||
return result
|
||||
create_mock_langsmith_client, should_use_langsmith_mock = create_mock_client_factory(_config)
|
||||
|
|
|
|||
191
litellm/integrations/mock_client_factory.py
Normal file
191
litellm/integrations/mock_client_factory.py
Normal file
|
|
@ -0,0 +1,191 @@
|
|||
"""
|
||||
Factory for creating mock HTTP clients for integration testing.
|
||||
|
||||
This module provides a simple factory pattern to create mock clients that intercept
|
||||
API calls and return successful mock responses, allowing full code execution without
|
||||
making actual network calls.
|
||||
"""
|
||||
|
||||
import httpx
|
||||
import json
|
||||
import asyncio
|
||||
from datetime import timedelta
|
||||
from typing import Dict, Optional, Callable, List, cast
|
||||
from dataclasses import dataclass
|
||||
|
||||
from litellm._logging import verbose_logger
|
||||
|
||||
|
||||
@dataclass
|
||||
class MockClientConfig:
|
||||
"""Configuration for creating a mock client."""
|
||||
name: str # e.g., "GCS", "LANGFUSE", "LANGSMITH", "DATADOG"
|
||||
env_var: str # e.g., "GCS_MOCK", "LANGFUSE_MOCK"
|
||||
default_latency_ms: int = 100 # Default mock latency in milliseconds
|
||||
default_status_code: int = 200 # Default HTTP status code
|
||||
default_json_data: Optional[Dict] = None # Default JSON response data
|
||||
url_matchers: Optional[List[str]] = None # List of strings to match in URLs (e.g., ["storage.googleapis.com"])
|
||||
patch_async_handler: bool = True # Whether to patch AsyncHTTPHandler.post
|
||||
patch_sync_client: bool = False # Whether to patch httpx.Client.post
|
||||
|
||||
def __post_init__(self):
|
||||
"""Ensure url_matchers is a list."""
|
||||
if self.url_matchers is None:
|
||||
self.url_matchers = []
|
||||
|
||||
|
||||
class MockResponse:
|
||||
"""Generic mock httpx.Response that satisfies API requirements."""
|
||||
|
||||
def __init__(self, status_code: int = 200, 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": "success"}
|
||||
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("")
|
||||
self.elapsed = timedelta(seconds=elapsed_seconds)
|
||||
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_url_match(url, matchers: List[str]) -> bool:
|
||||
"""Check if URL matches any of the provided matchers."""
|
||||
try:
|
||||
parsed_url = httpx.URL(url) if isinstance(url, str) else url
|
||||
url_str = str(parsed_url).lower()
|
||||
hostname = parsed_url.host or ""
|
||||
|
||||
for matcher in matchers:
|
||||
if matcher.lower() in url_str or matcher.lower() in hostname.lower():
|
||||
return True
|
||||
|
||||
# Also check for localhost with matcher in path
|
||||
if hostname in ("localhost", "127.0.0.1"):
|
||||
for matcher in matchers:
|
||||
if matcher.lower() in url_str:
|
||||
return True
|
||||
|
||||
return False
|
||||
except Exception:
|
||||
return False
|
||||
|
||||
|
||||
def create_mock_client_factory(config: MockClientConfig):
|
||||
"""
|
||||
Factory function that creates mock client functions based on configuration.
|
||||
|
||||
Returns:
|
||||
tuple: (create_mock_client_func, should_use_mock_func)
|
||||
"""
|
||||
# Store original methods for restoration
|
||||
_original_async_handler_post = None
|
||||
_original_sync_client_post = None
|
||||
_mocks_initialized = False
|
||||
|
||||
# Calculate mock latency
|
||||
import os
|
||||
latency_env = f"{config.name.upper()}_MOCK_LATENCY_MS"
|
||||
_MOCK_LATENCY_SECONDS = float(os.getenv(latency_env, str(config.default_latency_ms))) / 1000.0
|
||||
|
||||
# Create URL matcher function
|
||||
def _is_mock_url(url) -> bool:
|
||||
# url_matchers is guaranteed to be a list after __post_init__
|
||||
return _is_url_match(url, cast(List[str], config.url_matchers))
|
||||
|
||||
# Create async handler mock
|
||||
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 API calls."""
|
||||
if isinstance(url, str) and _is_mock_url(url):
|
||||
verbose_logger.info(f"[{config.name} MOCK] POST to {url}")
|
||||
await asyncio.sleep(_MOCK_LATENCY_SECONDS)
|
||||
return MockResponse(
|
||||
status_code=config.default_status_code,
|
||||
json_data=config.default_json_data,
|
||||
url=url,
|
||||
elapsed_seconds=_MOCK_LATENCY_SECONDS
|
||||
)
|
||||
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)
|
||||
raise RuntimeError("Original AsyncHTTPHandler.post not available")
|
||||
|
||||
# Create sync client mock
|
||||
def _mock_sync_client_post(self, url, **kwargs):
|
||||
"""Monkey-patched httpx.Client.post that intercepts API calls."""
|
||||
if _is_mock_url(url):
|
||||
verbose_logger.info(f"[{config.name} MOCK] POST to {url} (sync)")
|
||||
return MockResponse(
|
||||
status_code=config.default_status_code,
|
||||
json_data=config.default_json_data,
|
||||
url=url,
|
||||
elapsed_seconds=_MOCK_LATENCY_SECONDS
|
||||
)
|
||||
if _original_sync_client_post is not None:
|
||||
return _original_sync_client_post(self, url, **kwargs)
|
||||
|
||||
# Create mock client initialization function
|
||||
def create_mock_client():
|
||||
"""Initialize the mock client by patching HTTP handlers."""
|
||||
nonlocal _original_async_handler_post, _original_sync_client_post, _mocks_initialized
|
||||
|
||||
if _mocks_initialized:
|
||||
return
|
||||
|
||||
verbose_logger.debug(f"[{config.name} MOCK] Initializing {config.name} mock client...")
|
||||
|
||||
if config.patch_async_handler and _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(f"[{config.name} MOCK] Patched AsyncHTTPHandler.post")
|
||||
|
||||
if config.patch_sync_client and _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(f"[{config.name} MOCK] Patched httpx.Client.post")
|
||||
|
||||
verbose_logger.debug(f"[{config.name} MOCK] Mock latency set to {_MOCK_LATENCY_SECONDS*1000:.0f}ms")
|
||||
verbose_logger.debug(f"[{config.name} MOCK] {config.name} mock client initialization complete")
|
||||
|
||||
_mocks_initialized = True
|
||||
|
||||
# Create should_use_mock function
|
||||
def should_use_mock() -> bool:
|
||||
"""Determine if mock mode should be enabled."""
|
||||
import os
|
||||
from litellm.secret_managers.main import str_to_bool
|
||||
|
||||
mock_mode = os.getenv(config.env_var, "false")
|
||||
result = str_to_bool(mock_mode)
|
||||
result = bool(result) if result is not None else False
|
||||
|
||||
if result:
|
||||
verbose_logger.info(f"{config.name} Mock Mode: ENABLED - API calls will be mocked")
|
||||
|
||||
return result
|
||||
|
||||
return create_mock_client, should_use_mock
|
||||
Loading…
Add table
Reference in a new issue