diff --git a/litellm/integrations/datadog/datadog_mock_client.py b/litellm/integrations/datadog/datadog_mock_client.py index 80c1003510d..a0a760deb0b 100644 --- a/litellm/integrations/datadog/datadog_mock_client.py +++ b/litellm/integrations/datadog/datadog_mock_client.py @@ -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) diff --git a/litellm/integrations/gcs_bucket/gcs_bucket_mock_client.py b/litellm/integrations/gcs_bucket/gcs_bucket_mock_client.py index 6201dc343dc..de3e1739482 100644 --- a/litellm/integrations/gcs_bucket/gcs_bucket_mock_client.py +++ b/litellm/integrations/gcs_bucket/gcs_bucket_mock_client.py @@ -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 diff --git a/litellm/integrations/langfuse/langfuse_mock_client.py b/litellm/integrations/langfuse/langfuse_mock_client.py index 1dc739ea328..8ed6cff8d47 100644 --- a/litellm/integrations/langfuse/langfuse_mock_client.py +++ b/litellm/integrations/langfuse/langfuse_mock_client.py @@ -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 diff --git a/litellm/integrations/langsmith_mock_client.py b/litellm/integrations/langsmith_mock_client.py index a23d0b8bf94..ef602908231 100644 --- a/litellm/integrations/langsmith_mock_client.py +++ b/litellm/integrations/langsmith_mock_client.py @@ -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) diff --git a/litellm/integrations/mock_client_factory.py b/litellm/integrations/mock_client_factory.py new file mode 100644 index 00000000000..d788570b995 --- /dev/null +++ b/litellm/integrations/mock_client_factory.py @@ -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