lijiaoqiao/llm-gateway-competitors/litellm-wheel-src/litellm/integrations/braintrust_mock_client.py

"""
Mock HTTP client for Braintrust integration testing.

This module intercepts Braintrust API calls and returns successful mock responses,
allowing full code execution without making actual network calls.

Usage:
    Set BRAINTRUST_MOCK=true in environment variables or config to enable mock mode.
"""

import os
import time
from urllib.parse import urlparse

from litellm._logging import verbose_logger
from litellm.integrations.mock_client_factory import (
    MockClientConfig,
    MockResponse,
    create_mock_client_factory,
)

# Use factory for should_use_mock and MockResponse
# Braintrust uses both HTTPHandler (sync) and AsyncHTTPHandler (async)
# Braintrust needs endpoint-specific responses, so we use custom HTTPHandler.post patching
_config = MockClientConfig(
    "BRAINTRUST",
    "BRAINTRUST_MOCK",
    default_latency_ms=100,
    default_status_code=200,
    default_json_data={"id": "mock-project-id", "status": "success"},
    url_matchers=[
        ".braintrustdata.com",
        "braintrustdata.com",
        ".braintrust.dev",
        "braintrust.dev",
    ],
    patch_async_handler=True,  # Patch AsyncHTTPHandler.post for async calls
    patch_sync_client=False,  # HTTPHandler uses self.client.send(), not self.client.post()
    patch_http_handler=False,  # We use custom patching for endpoint-specific responses
)

# Get should_use_mock and create_mock_client from factory
# We need to call the factory's create_mock_client to patch AsyncHTTPHandler.post
(
    create_mock_braintrust_factory_client,
    should_use_braintrust_mock,
) = create_mock_client_factory(_config)

# Store original HTTPHandler.post method (Braintrust-specific for sync calls with custom logic)
_original_http_handler_post = None
_mocks_initialized = False

# Default mock latency in seconds
_MOCK_LATENCY_SECONDS = float(os.getenv("BRAINTRUST_MOCK_LATENCY_MS", "100")) / 1000.0


def _is_braintrust_url(url: str) -> bool:
    """Check if URL is a Braintrust API URL."""
    if not isinstance(url, str):
        return False

    parsed = urlparse(url)
    host = (parsed.hostname or "").lower()

    if not host:
        return False

    return (
        host == "braintrustdata.com"
        or host.endswith(".braintrustdata.com")
        or host == "braintrust.dev"
        or host.endswith(".braintrust.dev")
    )


def _mock_http_handler_post(
    self,
    url,
    data=None,
    json=None,
    params=None,
    headers=None,
    timeout=None,
    stream=False,
    files=None,
    content=None,
    logging_obj=None,
):
    """Monkey-patched HTTPHandler.post that intercepts Braintrust calls with endpoint-specific responses."""
    # Only mock Braintrust API calls
    if isinstance(url, str) and _is_braintrust_url(url):
        verbose_logger.info(f"[BRAINTRUST MOCK] POST to {url}")
        time.sleep(_MOCK_LATENCY_SECONDS)
        # Return appropriate mock response based on endpoint
        if "/project" in url:
            # Project creation/retrieval/register endpoint
            project_name = json.get("name", "litellm") if json else "litellm"
            mock_data = {"id": f"mock-project-id-{project_name}", "name": project_name}
        elif "/project_logs" in url:
            # Log insertion endpoint
            mock_data = {"status": "success"}
        else:
            mock_data = _config.default_json_data
        return MockResponse(
            status_code=_config.default_status_code,
            json_data=mock_data,
            url=url,
            elapsed_seconds=_MOCK_LATENCY_SECONDS,
        )
    if _original_http_handler_post is not None:
        return _original_http_handler_post(
            self,
            url=url,
            data=data,
            json=json,
            params=params,
            headers=headers,
            timeout=timeout,
            stream=stream,
            files=files,
            content=content,
            logging_obj=logging_obj,
        )
    raise RuntimeError("Original HTTPHandler.post not available")


def create_mock_braintrust_client():
    """
    Monkey-patch HTTPHandler.post to intercept Braintrust sync calls.

    Braintrust uses HTTPHandler for sync calls and AsyncHTTPHandler for async calls.
    HTTPHandler.post uses self.client.send(), not self.client.post(), so we need
    custom patching for sync (similar to Helicone).
    AsyncHTTPHandler.post is patched by the factory.

    We use custom patching instead of factory's patch_http_handler because we need
    endpoint-specific responses (different for /project vs /project_logs).

    This function is idempotent - it only initializes mocks once, even if called multiple times.
    """
    global _original_http_handler_post, _mocks_initialized

    if _mocks_initialized:
        return

    verbose_logger.debug("[BRAINTRUST MOCK] Initializing Braintrust mock client...")

    from litellm.llms.custom_httpx.http_handler import HTTPHandler

    if _original_http_handler_post is None:
        _original_http_handler_post = HTTPHandler.post
        HTTPHandler.post = _mock_http_handler_post  # type: ignore
        verbose_logger.debug("[BRAINTRUST MOCK] Patched HTTPHandler.post")

    # CRITICAL: Call the factory's initialization function to patch AsyncHTTPHandler.post
    # This is required for async calls to be mocked
    create_mock_braintrust_factory_client()

    verbose_logger.debug(
        f"[BRAINTRUST MOCK] Mock latency set to {_MOCK_LATENCY_SECONDS*1000:.0f}ms"
    )
    verbose_logger.debug(
        "[BRAINTRUST MOCK] Braintrust mock client initialization complete"
    )

    _mocks_initialized = True