litellm/tests/test_litellm/proxy/conftest.py

"""
Shared fixtures and helpers for proxy tests.

This module provides reusable utilities for creating proxy test clients
with database and Redis cache configuration.
"""

import asyncio
import os
import tempfile
from typing import Dict, Optional

import pytest
import yaml
from fastapi.testclient import TestClient

_PROXY_MODULE_GLOBALS_TO_ISOLATE = (
    "master_key",
    "prisma_client",
)


class StubClientNotConnectedError(Exception):
    pass


class DisconnectedPrisma:
    """Mimics prisma-client-py after disconnect(): ``is_connected()`` is False
    and the ``_engine`` property raises ``ClientNotConnectedError``."""

    def is_connected(self) -> bool:
        return False

    @property
    def _engine(self) -> None:
        raise StubClientNotConnectedError(
            "Client is not connected to the query engine, you must call `connect()` "
            "before attempting to query data."
        )


@pytest.fixture
def disconnected_prisma() -> DisconnectedPrisma:
    """A stand-in for a Prisma client wedged in the disconnected state."""
    return DisconnectedPrisma()


@pytest.fixture(autouse=True)
def _isolate_proxy_module_globals():
    """
    Snapshot and restore module-level globals on litellm.proxy.proxy_server
    that tests sometimes mutate via raw setattr (not monkeypatch).

    Without this, a leaked value — e.g. master_key set by a sibling test —
    flips the auth short-circuit in user_api_key_auth and causes unrelated
    tests in the same xdist worker to return 401 instead of 200.
    """
    from litellm.proxy import proxy_server

    sentinel = object()
    snapshot = {
        name: getattr(proxy_server, name, sentinel)
        for name in _PROXY_MODULE_GLOBALS_TO_ISOLATE
    }
    try:
        yield
    finally:
        for name, value in snapshot.items():
            if value is sentinel:
                if hasattr(proxy_server, name):
                    delattr(proxy_server, name)
            else:
                setattr(proxy_server, name, value)


@pytest.fixture(autouse=True)
def _reset_graceful_shutdown_state():
    """Graceful shutdown state is process-scoped; keep it from leaking between tests."""
    from litellm.proxy.shutdown.graceful_shutdown_manager import (
        GracefulShutdownManager,
    )

    GracefulShutdownManager.reset()
    yield
    GracefulShutdownManager.reset()


def build_cache_config(enable_cache: bool = True) -> Optional[Dict]:
    """
    Build Redis cache configuration from environment variables.

    Args:
        enable_cache: Whether to enable cache (default: True)

    Returns:
        dict: Cache configuration dict with 'cache' and 'cache_params' keys, or None
    """
    if not enable_cache:
        return None

    redis_host = os.getenv("REDIS_HOST")
    if not redis_host:
        return None

    redis_port = os.getenv("REDIS_PORT", "6379")
    cache_params = {
        "type": "redis",
        "host": redis_host,
        "port": int(redis_port) if redis_port.isdigit() else redis_port,
    }

    redis_password = os.getenv("REDIS_PASSWORD")
    if redis_password:
        cache_params["password"] = redis_password

    return {"cache": True, "cache_params": cache_params}


def build_minimal_proxy_config(
    database_url: Optional[str] = None, **init_options
) -> Dict:
    """
    Build a minimal proxy configuration YAML.

    Args:
        database_url: Optional database URL (falls back to DATABASE_URL env var)
        **init_options: Additional configuration options:
            - master_key: API key for authentication (default: "sk-1234")
            - enable_cache: Whether to enable Redis cache (default: True)
            - success_callback: Callback function for success events

    Returns:
        dict: Configuration dictionary ready to be written as YAML
    """
    config = {
        "general_settings": {"master_key": init_options.get("master_key", "sk-1234")},
        "litellm_settings": {},
    }

    # Configure database
    db_url = database_url or os.getenv("DATABASE_URL")
    if db_url:
        config["general_settings"]["database_url"] = db_url

    # Configure cache if Redis is available
    enable_cache = init_options.get("enable_cache", True)
    cache_config = build_cache_config(enable_cache=enable_cache)
    if cache_config:
        config["litellm_settings"].update(cache_config)

    # Add success_callback if provided (for realistic readiness endpoint)
    if init_options.get("success_callback") is not None:
        config["litellm_settings"]["success_callback"] = init_options[
            "success_callback"
        ]

    # Add any other litellm_settings from init_options
    excluded_keys = {
        "master_key",
        "debug",
        "success_callback",
        "database_url",
        "enable_cache",
    }
    for key, value in init_options.items():
        if key not in excluded_keys and key not in config["litellm_settings"]:
            config["litellm_settings"][key] = value

    return config


def set_proxy_environment_variables(
    monkeypatch, database_url: Optional[str] = None
) -> None:
    """
    Set environment variables for database and Redis.

    Args:
        monkeypatch: pytest monkeypatch fixture
        database_url: Optional database URL (falls back to DATABASE_URL env var)
    """
    # Set database URL
    db_url = database_url or os.getenv("DATABASE_URL")
    if db_url:
        monkeypatch.setenv("DATABASE_URL", db_url)

    # Set Redis environment variables if available
    redis_host = os.getenv("REDIS_HOST")
    if redis_host:
        monkeypatch.setenv("REDIS_HOST", redis_host)
        monkeypatch.setenv("REDIS_PORT", os.getenv("REDIS_PORT", "6379"))
        redis_password = os.getenv("REDIS_PASSWORD")
        if redis_password:
            monkeypatch.setenv("REDIS_PASSWORD", redis_password)


def create_proxy_test_client(
    monkeypatch, database_url: Optional[str] = None, **init_options
) -> TestClient:
    """
    Create a proxy TestClient with optional database and Redis cache configuration.

    Args:
        monkeypatch: pytest monkeypatch fixture
        database_url: Optional database URL (falls back to DATABASE_URL env var)
        **init_options: Additional configuration options:
            - master_key: API key for authentication (default: "sk-1234")
            - enable_cache: Whether to enable Redis cache (default: True)
            - success_callback: Callback function for success events
            - debug: Enable debug mode

    Returns:
        TestClient: FastAPI test client for the proxy server
    """
    from litellm.proxy.proxy_server import (
        cleanup_router_config_variables,
        initialize,
        app,
    )

    cleanup_router_config_variables()

    # Get config file path
    filepath = os.path.dirname(os.path.abspath(__file__))
    default_config_fp = os.path.join(
        filepath, "test_configs", "test_config_no_auth.yaml"
    )

    # Check if we need to create a minimal config with Redis/database
    enable_cache = init_options.get("enable_cache", True)
    needs_redis = enable_cache and os.getenv("REDIS_HOST") is not None
    needs_db = (database_url or os.getenv("DATABASE_URL")) is not None

    # Create minimal config if:
    # 1. Default config file doesn't exist, OR
    # 2. We need Redis/database config that might not be in the default config
    if not os.path.exists(default_config_fp) or needs_redis or needs_db:
        minimal_config = build_minimal_proxy_config(
            database_url=database_url, **init_options
        )

        with tempfile.NamedTemporaryFile(mode="w", suffix=".yaml", delete=False) as f:
            yaml.dump(minimal_config, f)
            config_fp = f.name
    else:
        config_fp = default_config_fp

    # Set environment variables
    set_proxy_environment_variables(monkeypatch, database_url=database_url)

    # Initialize proxy
    asyncio.run(initialize(config=config_fp, debug=init_options.get("debug", False)))
    return TestClient(app)