litellm/tests/unit/proxy/conftest.py

# conftest.py

import asyncio
import copy
import inspect
import warnings
from collections.abc import Iterator
from typing import Dict

import pytest


import litellm
import litellm.proxy.proxy_server
from tests.unit.litellm_core_utils.fake_secret_vault import FakeSecretVault


# Top-level assignments of these types are the ones importlib.reload(litellm)
# would have effectively reset. We snapshot them at conftest import time and
# deep-copy the snapshot back before every test.
_SNAPSHOT_TYPES = (list, dict, set, tuple, str, int, float, bool, bytes)


def _snapshot_mutable_state(module):
    """Capture a per-module snapshot of primitive and collection attributes."""
    snapshot = {}
    for attr in list(vars(module)):
        if attr.startswith("_"):
            continue
        try:
            value = getattr(module, attr)
        except Exception as exc:
            warnings.warn(
                f"conftest: could not read {module.__name__}.{attr} during snapshot: {exc}",
                stacklevel=2,
            )
            continue
        if value is None or isinstance(value, _SNAPSHOT_TYPES):
            try:
                snapshot[attr] = copy.deepcopy(value)
            except Exception as exc:
                warnings.warn(
                    f"conftest: could not snapshot {module.__name__}.{attr}: {exc}",
                    stacklevel=2,
                )
    return snapshot


def _restore_mutable_state(module, snapshot):
    for attr, default in snapshot.items():
        try:
            setattr(module, attr, copy.deepcopy(default))
        except Exception as exc:
            warnings.warn(
                f"conftest: could not restore {module.__name__}.{attr}: {exc}",
                stacklevel=2,
            )


def _collect_flushable_caches():
    """Return (module, attr) pairs whose values expose flush_cache()."""
    targets = []
    for module in (litellm, litellm.proxy.proxy_server):
        for attr in list(vars(module)):
            if attr.startswith("_"):
                continue
            try:
                value = getattr(module, attr)
            except Exception:
                continue
            # Only instances — a class reference has an unbound flush_cache
            # that can't be called without a self argument.
            if inspect.isclass(value) or inspect.ismodule(value):
                continue
            if callable(getattr(value, "flush_cache", None)):
                targets.append((module, attr))
    return targets


def _flush_caches(targets):
    for module, attr in targets:
        try:
            value = getattr(module, attr)
        except Exception:
            continue
        flush = getattr(value, "flush_cache", None)
        if callable(flush):
            try:
                flush()
            except Exception as exc:
                warnings.warn(
                    f"conftest: flush_cache failed on {module.__name__}.{attr}: {exc}",
                    stacklevel=2,
                )


# Snapshot once at conftest import — these are the "clean" module states.
_LITELLM_STATE = _snapshot_mutable_state(litellm)
_PROXY_SERVER_STATE = _snapshot_mutable_state(litellm.proxy.proxy_server)
_FLUSHABLE_CACHES = _collect_flushable_caches()


@pytest.fixture(scope="function", autouse=True)
def setup_and_teardown():
    """Reset mutable module state on litellm and proxy_server before each test.

    Replaces a previous importlib.reload(litellm) approach that cost ~17s
    per test (re-executing the full litellm __init__ import chain).

    What IS reset:
      - Top-level module attributes of type list / dict / set / tuple
        / str / int / float / bool / bytes, and None-valued attributes.
        These cover callback lists, general_settings, master_key,
        premium_user, prisma_client, etc. — anything the old reload() reset
        by re-executing the module body.
      - Any module-level object instance that exposes flush_cache() (the
        DualCache and LLMClientCache family), which handles cache state
        that can't round-trip through deepcopy because of internal locks.

    What is NOT reset:
      - Class instances without flush_cache() (e.g. ProxyLogging,
        JWTHandler, FastAPI routers, loggers). If a test mutates such an
        instance in-place (setattr on the instance, appending to one of
        its internal lists, etc.), the mutation will leak into later tests.
        Use pytest's monkeypatch.setattr() or a local fixture for those
        cases — don't rely on this autouse fixture to undo them.
    """
    _restore_mutable_state(litellm, _LITELLM_STATE)
    _restore_mutable_state(litellm.proxy.proxy_server, _PROXY_SERVER_STATE)
    _flush_caches(_FLUSHABLE_CACHES)

    loop = asyncio.get_event_loop_policy().new_event_loop()
    asyncio.set_event_loop(loop)
    try:
        yield
    finally:
        loop.close()
        asyncio.set_event_loop(None)


def pytest_collection_modifyitems(config, items):
    # Separate tests in 'test_amazing_proxy_custom_logger.py' and other tests
    custom_logger_tests = [
        item for item in items if "custom_logger" in item.parent.name
    ]
    other_tests = [item for item in items if "custom_logger" not in item.parent.name]

    # Sort tests based on their names
    custom_logger_tests.sort(key=lambda x: x.name)
    other_tests.sort(key=lambda x: x.name)

    # Reorder the items list
    items[:] = custom_logger_tests + other_tests


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

_proxy_module_globals_snapshot = pytest.StashKey[Dict[str, object]]()


@pytest.hookimpl(hookwrapper=True)
def pytest_runtest_setup(item):
    from litellm.proxy import proxy_server

    item.stash[_proxy_module_globals_snapshot] = {
        name: vars(proxy_server)[name]
        for name in _PROXY_MODULE_GLOBALS_TO_ISOLATE
        if name in vars(proxy_server)
    }
    yield


@pytest.hookimpl(hookwrapper=True)
def pytest_runtest_teardown(item, nextitem):
    yield
    snapshot = item.stash.get(_proxy_module_globals_snapshot, None)
    if snapshot is None:
        return
    from litellm.proxy import proxy_server

    for name in _PROXY_MODULE_GLOBALS_TO_ISOLATE:
        if name in snapshot:
            setattr(proxy_server, name, snapshot[name])
        elif name in vars(proxy_server):
            delattr(proxy_server, name)


@pytest.fixture
def secret_vault_factory() -> type[FakeSecretVault]:
    return FakeSecretVault


@pytest.fixture
def httpx_transport(monkeypatch: pytest.MonkeyPatch) -> Iterator[None]:
    monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
    litellm.in_memory_llm_clients_cache.flush_cache()
    yield
    litellm.in_memory_llm_clients_cache.flush_cache()


@pytest.fixture(autouse=True)
def _reset_graceful_shutdown_state():
    from litellm.proxy.shutdown.graceful_shutdown_manager import (
        GracefulShutdownManager,
    )

    GracefulShutdownManager.reset()
    yield
    GracefulShutdownManager.reset()