litellm/tests/unit/conftest.py

import asyncio
import base64
import importlib
import os
import shutil
import tempfile
from collections.abc import AsyncIterator, Callable, Container, Coroutine, Iterable, Iterator
from dataclasses import dataclass, field
from pathlib import Path
from types import MappingProxyType
from typing import TYPE_CHECKING, Final, Literal, TypeAlias

import boto3
import httpx
import pytest
from pytest_socket import enable_socket, socket_allow_hosts

if TYPE_CHECKING:
    from unittest.mock import MagicMock

    from litellm.caching.redis_cache import RedisCache

RouterRedis: TypeAlias = tuple["RedisCache", asyncio.Queue[tuple[tuple[object, ...], str | None, str | None]]]
RedisWriteGate: TypeAlias = tuple[asyncio.Event, asyncio.Event, asyncio.Event]


def redis_cache_double() -> "MagicMock":
    """A `RedisCache` double for router tests. The native router stores connect from its settings, which name
    a closed local port, so they serve their memory tier while the double answers the Python Redis calls."""
    from unittest.mock import MagicMock

    from litellm.caching.redis_cache import RedisCache
    from litellm.redis.store_settings import RedisStoreSettings

    redis: Final = MagicMock(spec=RedisCache)
    redis.redis_kwargs = {"host": "127.0.0.1", "port": 1, "socket_timeout": 0.2}
    redis.namespace = None
    redis.store_settings = RedisStoreSettings(
        redis_kwargs=redis.redis_kwargs, namespace=None, default_ttl=60, breaker=None
    )
    return redis


@pytest.fixture
def redis_write_gate() -> RedisWriteGate:
    released: Final = asyncio.Event()
    released.set()
    return asyncio.Event(), released, asyncio.Event()


@pytest.fixture
async def router_redis(
    redis_write_gate: RedisWriteGate,
) -> AsyncIterator[RouterRedis]:
    from unittest.mock import patch

    from fakeredis import FakeRedis, FakeServer
    from fakeredis.aioredis import FakeConnection
    from opentelemetry import trace
    from redis.asyncio import BlockingConnectionPool
    from redis.exceptions import ConnectionError

    from litellm._internal_context import current_service_target
    from litellm.caching.redis_cache import RedisCache

    commands: Final[asyncio.Queue[tuple[tuple[object, ...], str | None, str | None]]] = asyncio.Queue()
    started, released, unavailable = redis_write_gate

    class ObservedConnection(FakeConnection):
        async def send_command(self, *args: object, **kwargs: object) -> None:  # kwargs-ok: Redis transport signature
            commands.put_nowait((args, current_service_target(), getattr(trace.get_current_span(), "name", None)))
            await super().send_command(*args, **kwargs)

        async def send_packed_command(self, command: Iterable[bytes] | bytes | str, check_health: bool = True) -> None:
            chunks: Final = (
                (command.encode(),)
                if isinstance(command, str)
                else (command,)
                if isinstance(command, bytes)
                else command
            )
            wire: Final = b"".join(chunks)
            if b"INCRBYFLOAT" in wire:
                commands.put_nowait(
                    ((wire,), current_service_target(), getattr(trace.get_current_span(), "name", None))
                )
                started.set()
                await released.wait()
                if unavailable.is_set():
                    raise ConnectionError("scripted write outage")
            await super().send_packed_command(command, check_health=check_health)

    server: Final = FakeServer()
    client: Final = FakeRedis(server=server)
    pool: Final = BlockingConnectionPool(connection_class=ObservedConnection, server=server)
    try:
        with (
            patch("redis.Redis.from_url", return_value=client),
            patch("redis.asyncio.BlockingConnectionPool.from_url", return_value=pool),
        ):
            yield RedisCache(url="redis://router-unit.example:6379/0"), commands
    finally:
        released.set()
        await pool.aclose()
        client.close()


HOST_ENVIRONMENT_ALLOWLIST: Final = frozenset(
    (
        "PATH",
        "HOME",
        "USER",
        "LOGNAME",
        "TMPDIR",
        "TEMP",
        "TMP",
        "LANG",
        "LC_ALL",
        "LC_CTYPE",
        "TZ",
        "VIRTUAL_ENV",
        "LITELLM_LOCAL_MODEL_COST_MAP",
        "TIKTOKEN_CACHE_DIR",
    )
)
HOST_ENVIRONMENT_ALLOWED_PREFIXES: Final = ("PYTEST_", "PYTHON", "COV_CORE_", "COVERAGE_")
HOST_ONLY_ENVIRONMENT: Final = frozenset(
    name
    for name in os.environ
    if name not in HOST_ENVIRONMENT_ALLOWLIST and not name.startswith(HOST_ENVIRONMENT_ALLOWED_PREFIXES)
)

os.environ["PYTHON_DOTENV_DISABLED"] = "1"
os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True"

import litellm  # noqa: E402  # litellm reads LITELLM_LOCAL_MODEL_COST_MAP at import
from litellm.litellm_core_utils.model_cost_table import publish  # noqa: E402  # same import-time dependency
from litellm._logging import ALL_LOGGERS  # noqa: E402  # same import-time dependency
from litellm.anthropic_beta_headers_manager import (
    reload_beta_headers_config,  # noqa: E402  # same import-time dependency
)
from litellm.litellm_core_utils.get_model_cost_map import (  # noqa: E402  # same import-time dependency
    model_cost_inputs,
)
from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER  # noqa: E402  # same import-time dependency
from litellm.litellm_core_utils.prompt_templates import (
    factory as prompt_factory_module,  # noqa: E402  # same import-time dependency
)
from litellm.litellm_core_utils.prompt_templates import (  # noqa: E402  # same import-time dependency
    image_handling as image_handling_module,
)
from litellm.llms.custom_httpx.async_client_cleanup import (  # noqa: E402  # same import-time dependency
    close_litellm_async_clients,
)
from litellm.llms.gemini.chat import (
    transformation as gemini_chat_transformation_module,  # noqa: E402  # same import-time dependency
)
from litellm.proxy.db import (
    tool_registry_writer as tool_registry_writer_module,  # noqa: E402  # same import-time dependency
)
from litellm.router_rust.surface.replay import (  # noqa: E402  # same import-time dependency
    _live_registries as live_router_registries,  # pyright: ignore[reportPrivateUsage]  # the routers the model-cost replay reaches
)
from litellm.router_utils.global_callback_holders import (  # noqa: E402  # same import-time dependency
    routers_holding_global_callbacks,
)

LOOPBACK_HOSTS: Final = ["127.0.0.1", "::1", "localhost"]
AMBIENT_AZURE_CREDENTIAL_ENV_VARS: Final = (
    "AZURE_AD_TOKEN",
    "AZURE_TENANT_ID",
    "AZURE_CLIENT_ID",
    "AZURE_CLIENT_SECRET",
    "AZURE_USERNAME",
    "AZURE_PASSWORD",
)
AMBIENT_AWS_ENV_VARS: Final = (
    "AWS_PROFILE",
    "AWS_DEFAULT_PROFILE",
    "AWS_CONTAINER_CREDENTIALS_FULL_URI",
    "AWS_CONTAINER_CREDENTIALS_RELATIVE_URI",
    "AWS_SESSION_TOKEN",
    "AWS_ROLE_ARN",
    "AWS_WEB_IDENTITY_TOKEN_FILE",
    "AWS_BEARER_TOKEN_BEDROCK",
    "AWS_REGION_NAME",
    "AWS_DEFAULT_REGION",
)
MODULES_WITH_AWS_AUTH_HANDLERS: Final = (
    "litellm.main",
    "litellm.files.main",
    "litellm.rerank_api.main",
    "litellm.realtime_api.main",
)
CALLBACK_LISTS: Final = (
    "callbacks",
    "success_callback",
    "failure_callback",
    "input_callback",
    "_async_success_callback",
    "_async_failure_callback",
    "_async_input_callback",
)
RESET_TO_NONE_GLOBALS: Final = (
    "model_fallbacks",
    "cache",
    "max_fallbacks",
    "default_fallbacks",
    "fallbacks",
    "context_window_fallbacks",
    "content_policy_fallbacks",
)
RESTORED_GLOBALS: Final = (
    "disable_aiohttp_transport",
    "force_ipv4",
    "drop_params",
    "secret_manager_client",
    "_key_management_system",
    "_key_management_settings",
    "api_base",
    "num_retries",
    "modify_params",
    "ssl_verify",
    "credential_list",
    "model_group_settings",
    "default_internal_user_params",
    "default_team_params",
    "prometheus_emit_stream_label",
    "vector_store_registry",
    "model_cost",
    "cost_margin_config",
    "cost_discount_config",
    "disable_hf_tokenizer_download",
    "disable_copilot_system_to_assistant",
    "cohere_models",
    "anthropic_models",
    "token_counter",
    "initialized_langfuse_clients",
)
MODULE_LEVEL_CLIENTS: Final = ("module_level_client", "module_level_aclient")
SESSION_CLIENTS: Final = ("base_llm_aiohttp_handler", "httpx_client", "aclient", "client")
ONE_PIXEL_PNG: Final = base64.b64decode(
    "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNkYPhfDwAChwGA60e6kgAAAABJRU5ErkJggg=="
)


def _allow_loopback_only() -> None:
    socket_allow_hosts(LOOPBACK_HOSTS, allow_unix_socket=True)


_allow_loopback_only()


def pytest_collectstart() -> None:
    _allow_loopback_only()


@pytest.hookimpl(trylast=True)
def pytest_runtest_setup() -> None:
    _allow_loopback_only()


def _run_coroutine_if_needed(result: object) -> None:
    if not asyncio.iscoroutine(result):
        return
    coroutine: Final[Coroutine[object, object, object]] = result
    try:
        asyncio.run(coroutine)
    except RuntimeError:
        try:
            loop: Final = asyncio.get_running_loop()
        except RuntimeError:
            coroutine.close()
            return
        loop.create_task(coroutine)


async def _complete_test_logging() -> None:
    await GLOBAL_LOGGING_WORKER.flush()
    await GLOBAL_LOGGING_WORKER.stop()


def _flush_completed_test_logging() -> None:
    queue: Final = GLOBAL_LOGGING_WORKER._queue
    if queue is None or queue._unfinished_tasks == 0:
        return
    worker_loop: Final = GLOBAL_LOGGING_WORKER._bound_loop
    if worker_loop is None or worker_loop.is_closed():
        _run_coroutine_if_needed(_complete_test_logging())
        return
    if worker_loop.is_running():
        asyncio.run_coroutine_threadsafe(_complete_test_logging(), worker_loop).result()
        return
    worker_loop.run_until_complete(_complete_test_logging())


@pytest.hookimpl(hookwrapper=True)
def pytest_runtest_call() -> Iterator[None]:
    yield
    _flush_completed_test_logging()


def _close_handler_if_needed(handler: object) -> None:
    close: Final = getattr(handler, "close", None)
    if not callable(close):
        return
    _run_coroutine_if_needed(close())


def _reset_aws_auth_caches() -> None:
    modules: Final = tuple(importlib.import_module(name) for name in MODULES_WITH_AWS_AUTH_HANDLERS)
    flushes: Final = (
        getattr(getattr(getattr(module, attr_name), "iam_cache", None), "flush_cache", None)
        for module in modules
        for attr_name in dir(module)
    )
    for flush in filter(callable, flushes):
        flush()
    boto3.DEFAULT_SESSION = None


def _flush_client_caches() -> None:
    litellm.in_memory_llm_clients_cache.flush_cache()
    image_handling_module.in_memory_cache.flush_cache()
    image_handling_module.pdf_probe_verdicts.flush_cache()
    _reset_aws_auth_caches()


_MACOS_LONGEST_SOCKET_PATH: Final = 103
_LONGEST_SOCKET_SUFFIX: Final = len("/sock12345678/d/litellm-pgbouncer-12345678/.s.PGSQL.65535")


def _socket_root() -> str:
    default: Final = tempfile.gettempdir()
    fits: Final = len(os.fsencode(default)) + _LONGEST_SOCKET_SUFFIX <= _MACOS_LONGEST_SOCKET_PATH
    return default if fits else "/tmp"


@pytest.fixture
def socket_dir(tmp_path: Path) -> Iterator[Path]:
    """``tmp_path`` behind a short link, so unix socket paths fit the sun_path limit that it exceeds on macOS."""
    link_dir: Final = Path(tempfile.mkdtemp(prefix="sock", dir=_socket_root()))
    link: Final = link_dir / "d"
    link.symlink_to(tmp_path, target_is_directory=True)
    yield link
    shutil.rmtree(link_dir, ignore_errors=True)


@pytest.fixture
def socket_path(socket_dir: Path) -> Path:
    return socket_dir / "spend.sock"


@pytest.fixture
def socket_tempdir(socket_dir: Path, monkeypatch: pytest.MonkeyPatch) -> Path:
    """Point ``tempfile``'s default directory at ``socket_dir``, for code that picks its own socket directory."""
    monkeypatch.setattr(tempfile, "tempdir", str(socket_dir))
    return socket_dir


@pytest.fixture(autouse=True, scope="session")
def bundled_tiktoken_cache() -> None:
    importlib.import_module("litellm.litellm_core_utils.default_encoding")


@pytest.fixture(scope="session")
def isolated_aws_config_files(tmp_path_factory: pytest.TempPathFactory) -> tuple[Path, Path]:
    aws_dir: Final = tmp_path_factory.mktemp("aws-config")
    credentials: Final = aws_dir / "credentials"
    config: Final = aws_dir / "config"
    credentials.write_text("", encoding="utf-8")
    config.write_text("", encoding="utf-8")
    return credentials, config


@pytest.fixture(autouse=True)
def isolate_host_environment(isolated_aws_config_files: tuple[Path, Path]) -> Iterator[None]:
    credentials, config = isolated_aws_config_files
    with pytest.MonkeyPatch.context() as environment:
        for name in HOST_ONLY_ENVIRONMENT:
            environment.delenv(name, raising=False)
        environment.setenv("AWS_SHARED_CREDENTIALS_FILE", str(credentials))
        environment.setenv("AWS_CONFIG_FILE", str(config))
        environment.setenv("AWS_EC2_METADATA_DISABLED", "true")
        for name in AMBIENT_AWS_ENV_VARS:
            environment.delenv(name, raising=False)
        for name in ("PROXY_BASE_URL", "SERVER_ROOT_PATH", "SERVER_ROOT_PATHS", "LITELLM_LOG"):
            environment.delenv(name, raising=False)
        environment.setenv("LITELLM_CLI_DISABLE_KEYRING", "1")
        yield


@pytest.fixture(autouse=True)
def isolate_litellm_globals() -> Iterator[None]:
    original_callbacks: Final = {name: list(getattr(litellm, name) or []) for name in CALLBACK_LISTS}
    original_callback_holders: Final = frozenset(routers_holding_global_callbacks)
    original_reset: Final = {name: getattr(litellm, name) for name in RESET_TO_NONE_GLOBALS}
    original_restored: Final = {name: getattr(litellm, name) for name in RESTORED_GLOBALS if hasattr(litellm, name)}
    original_clients: Final = {
        name: litellm.__dict__[name] for name in MODULE_LEVEL_CLIENTS if name in litellm.__dict__
    }
    original_loggers: Final = {
        logger: (logger.level, logger.disabled, logger.propagate, list(logger.handlers), list(logger.filters))
        for logger in ALL_LOGGERS
    }
    original_tool_policy_registry: Final = tool_registry_writer_module._tool_policy_registry
    _flush_client_caches()
    for name in CALLBACK_LISTS:
        setattr(litellm, name, [])
    routers_holding_global_callbacks.clear()
    for name in RESET_TO_NONE_GLOBALS:
        setattr(litellm, name, None)
    for name in MODULE_LEVEL_CLIENTS:
        litellm.__dict__.pop(name, None)
    tool_registry_writer_module._tool_policy_registry = None
    yield
    _flush_client_caches()
    leaked_clients: Final = tuple(litellm.__dict__.pop(name, None) for name in MODULE_LEVEL_CLIENTS)
    for name, client in zip(MODULE_LEVEL_CLIENTS, leaked_clients):
        if client is not original_clients.get(name):
            _close_handler_if_needed(client)
    litellm.__dict__.update(original_clients)
    for name, value in (original_callbacks | original_reset | original_restored).items():
        setattr(litellm, name, value)
    routers_holding_global_callbacks.clear()
    routers_holding_global_callbacks.update(original_callback_holders)
    for logger, (level, disabled, propagate, handlers, filters) in original_loggers.items():
        logger.setLevel(level)
        logger.disabled = disabled
        logger.propagate = propagate
        logger.handlers = handlers
        logger.filters = filters
    tool_registry_writer_module._tool_policy_registry = original_tool_policy_registry


_IMPORTED_MODEL_COST_INPUTS: Final = model_cost_inputs.snapshot()


@pytest.fixture(autouse=True)
def isolate_router_model_cost_state() -> Iterator[None]:
    original_live_registries: Final = frozenset(live_router_registries)
    original_model_cost_inputs: Final = model_cost_inputs.snapshot()
    model_cost_inputs.restore(_IMPORTED_MODEL_COST_INPUTS)
    publish(litellm.model_cost)
    live_router_registries.clear()
    yield
    live_router_registries.clear()
    live_router_registries.update(original_live_registries)
    model_cost_inputs.restore(original_model_cost_inputs)
    publish(litellm.model_cost)
    litellm.get_model_info.cache_clear()


def _provider_model_sets() -> tuple[set[str], ...]:
    module_globals: Final = tuple(vars(litellm).items())
    return tuple(models for name, models in module_globals if name.endswith("_models") and isinstance(models, set))


@pytest.fixture(autouse=True)
def isolate_model_catalog() -> Iterator[None]:
    model_cost: Final = litellm.model_cost
    model_cost_entries: Final = dict(model_cost)
    models_by_provider: Final = litellm.models_by_provider
    models_by_provider_entries: Final = dict(models_by_provider)
    provider_model_sets: Final = tuple((models, frozenset(models)) for models in _provider_model_sets())
    yield
    model_cost.clear()
    model_cost.update(model_cost_entries)
    models_by_provider.clear()
    models_by_provider.update(models_by_provider_entries)
    for models, entries in provider_model_sets:
        models.clear()
        models.update(entries)
    publish(litellm.model_cost)
    litellm.get_model_info.cache_clear()


@pytest.fixture
def local_model_cost_map(monkeypatch: pytest.MonkeyPatch) -> Iterator[None]:
    monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True")
    monkeypatch.setattr(litellm, "model_cost", litellm.get_model_cost_map(url=""))
    litellm.get_model_info.cache_clear()
    yield
    litellm.get_model_info.cache_clear()


@pytest.fixture
def declare_mid_conversation_system(monkeypatch: pytest.MonkeyPatch) -> Callable[[str, bool], None]:
    def declare(model_cost_key: str, supported: bool) -> None:
        entry: Final = {
            "litellm_provider": model_cost_key.split("/", 1)[0],
            "mode": "chat",
            "supports_mid_conversation_system": supported,
        }
        monkeypatch.setitem(litellm.model_cost, model_cost_key, entry)
        litellm.get_model_info.cache_clear()

    return declare


NativeStructuredOutputKey: TypeAlias = Literal["legacy", "bedrock", "default"] | None


@pytest.fixture
def set_native_structured_output(
    local_model_cost_map: None, monkeypatch: pytest.MonkeyPatch
) -> Callable[[str, NativeStructuredOutputKey], None]:
    def set_flags(model_cost_key: str, native_key: NativeStructuredOutputKey) -> None:
        entry: Final = litellm.model_cost[model_cost_key]
        monkeypatch.setitem(entry, "supports_native_structured_output", native_key == "legacy")
        monkeypatch.setitem(entry, "supports_bedrock_native_structured_output", native_key in ("bedrock", "default"))
        monkeypatch.setitem(entry, "supports_bedrock_native_structured_output_by_default", native_key == "default")
        litellm.get_model_info.cache_clear()

    return set_flags


@pytest.fixture
def local_beta_headers_config(monkeypatch: pytest.MonkeyPatch) -> Iterator[None]:
    monkeypatch.setenv("LITELLM_LOCAL_ANTHROPIC_BETA_HEADERS", "True")
    reload_beta_headers_config()
    yield
    monkeypatch.delenv("LITELLM_LOCAL_ANTHROPIC_BETA_HEADERS", raising=False)
    reload_beta_headers_config()


@dataclass(slots=True)
class AsyncOnlyImageFetch:
    fetched: list[str] = field(default_factory=list)  # mutable-ok: tests assert on the URLs fetched, in order
    base64_png: str = base64.b64encode(ONE_PIXEL_PNG).decode()
    data_url: str = "data:image/png;base64," + base64.b64encode(ONE_PIXEL_PNG).decode()


@pytest.fixture
def async_only_image_fetch(monkeypatch: pytest.MonkeyPatch) -> AsyncOnlyImageFetch:
    fetch: Final = AsyncOnlyImageFetch()

    def forbid_sync_fetch(client: object, url: str, **kwargs: object) -> httpx.Response:
        raise litellm.ImageFetchError(f"sync image fetch ran on the event loop: {url}")

    async def serve_png(client: object, url: str, **kwargs: object) -> httpx.Response:
        fetch.fetched.append(url)
        return httpx.Response(
            200, content=ONE_PIXEL_PNG, headers={"content-type": "image/png"}, request=httpx.Request("GET", url)
        )

    def forbid_sync_convert(url: str, *args: object, **kwargs: object) -> str:
        if url.startswith(("http://", "https://")):
            raise litellm.ImageFetchError(f"sync convert_url_to_base64 ran on the request path: {url}")
        return url

    monkeypatch.setattr(image_handling_module, "safe_get", forbid_sync_fetch)
    monkeypatch.setattr(image_handling_module, "async_safe_get", serve_png)
    for module in (image_handling_module, prompt_factory_module, gemini_chat_transformation_module):
        monkeypatch.setattr(module, "convert_url_to_base64", forbid_sync_convert)
    return fetch


@pytest.fixture
def no_ambient_azure_credentials(monkeypatch: pytest.MonkeyPatch) -> None:
    for name in AMBIENT_AZURE_CREDENTIAL_ENV_VARS:
        monkeypatch.delenv(name, raising=False)


FAKE_PROVIDER_CREDENTIALS: Final = MappingProxyType(
    {
        "OPENAI_API_KEY": "sk-unit-test",
        "ANTHROPIC_API_KEY": "sk-ant-unit-test",
        "GEMINI_API_KEY": "unit-test",
        "AZURE_API_KEY": "unit-test",
        "AZURE_API_BASE": "https://unit-test.openai.azure.com",
        "AZURE_API_VERSION": "2024-02-01",
        "AWS_ACCESS_KEY_ID": "unit-test",
        "AWS_SECRET_ACCESS_KEY": "unit-test",
        "AWS_REGION_NAME": "us-east-1",
        "COHERE_API_KEY": "unit-test",
        "DD_API_KEY": "unit-test",
        "DD_SITE": "us5.datadoghq.com",
    }
)


@pytest.fixture
def fake_provider_credentials(monkeypatch: pytest.MonkeyPatch) -> None:
    for name, value in FAKE_PROVIDER_CREDENTIALS.items():
        monkeypatch.setenv(name, value)


@pytest.fixture
async def drained_logging_worker() -> None:
    await asyncio.wait_for(GLOBAL_LOGGING_WORKER.clear_queue(), timeout=10)


def pytest_sessionfinish() -> None:
    for name in MODULE_LEVEL_CLIENTS:
        _close_handler_if_needed(litellm.__dict__.pop(name, None))
    for name in SESSION_CLIENTS:
        _close_handler_if_needed(getattr(litellm, name, None))
    _run_coroutine_if_needed(close_litellm_async_clients())
    enable_socket()


def in_operator_error(value: Container[object]) -> str:
    with pytest.raises(TypeError) as native:
        _ = "key" in value
    return str(native.value)