litellm/tests/unified_google_tests/conftest.py

# conftest.py

import asyncio
import importlib
import os
import socket
import threading
import time
from pathlib import Path
from typing import Iterator, Tuple

import pytest
import uvicorn
from dotenv import load_dotenv

load_dotenv()

import litellm  # noqa: E402,F401

from tests._vcr_conftest_common import (  # noqa: E402,F401
    VerboseReporterState,
    _pin_multipart_boundary,
    apply_vcr_auto_marker_to_items,
    emit_cassette_cache_session_banner,
    emit_vcr_classification_summary,
    emit_vcr_diagnostic_log,
    install_live_call_probe,
    record_vcr_outcome,
    register_persister_if_enabled,
    reset_vcr_diag_dir,
    vcr_config_dict,
)

_verbose_state = VerboseReporterState()

PROXY_CONFIG_PATH = Path(__file__).parent / "google_genai_proxy_test_config.yaml"
PROXY_MASTER_KEY = "sk-1234"
PROXY_START_TIMEOUT_S = 30.0


def _start_proxy_server(
    config_path: str,
) -> Tuple[str, uvicorn.Server, threading.Thread, socket.socket]:
    from litellm.proxy.proxy_server import (
        app as proxy_app,
        cleanup_router_config_variables,
        initialize,
    )

    cleanup_router_config_variables()

    sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
    sock.bind(("127.0.0.1", 0))
    host, port = sock.getsockname()

    config = uvicorn.Config(proxy_app, host=host, port=port, log_level="warning")
    server = uvicorn.Server(config)

    def _run() -> None:
        loop = asyncio.new_event_loop()
        asyncio.set_event_loop(loop)
        loop.run_until_complete(initialize(config=config_path, debug=True))
        loop.run_until_complete(server.serve(sockets=[sock]))

    thread = threading.Thread(target=_run, daemon=True)
    thread.start()

    start_time = time.time()
    while not server.started:
        if not thread.is_alive():
            raise RuntimeError("LiteLLM proxy failed to start")
        if time.time() - start_time > PROXY_START_TIMEOUT_S:
            raise TimeoutError("LiteLLM proxy did not start in time")
        time.sleep(0.05)

    return f"http://{host}:{port}", server, thread, sock


@pytest.fixture(scope="session")
def google_genai_proxy_url() -> Iterator[str]:
    from base_google_genai_proxy_sdk_test import has_vertex_credentials
    from base_google_test import load_vertex_ai_credentials

    saved_env = {
        key: os.environ.get(key)
        for key in (
            "DATABASE_URL",
            "DIRECT_URL",
            "LITELLM_MASTER_KEY",
            "STORE_MODEL_IN_DB",
            "GOOGLE_APPLICATION_CREDENTIALS",
        )
    }
    temp_credentials_path: str | None = None
    os.environ.pop("DATABASE_URL", None)
    os.environ.pop("DIRECT_URL", None)
    os.environ["LITELLM_MASTER_KEY"] = PROXY_MASTER_KEY
    os.environ["STORE_MODEL_IN_DB"] = "False"

    if has_vertex_credentials():
        credentials_file = os.environ.get("GOOGLE_APPLICATION_CREDENTIALS", "")
        if not (credentials_file and os.path.isfile(credentials_file)):
            vertex_credentials_path = load_vertex_ai_credentials(
                model="vertex_ai/gemini-2.5-flash-lite"
            )
            if vertex_credentials_path:
                temp_credentials_path = vertex_credentials_path
                os.environ["GOOGLE_APPLICATION_CREDENTIALS"] = vertex_credentials_path

    server_url, server, thread, sock = _start_proxy_server(str(PROXY_CONFIG_PATH))
    try:
        yield server_url
    finally:
        server.should_exit = True
        thread.join(timeout=10)
        sock.close()
        if temp_credentials_path:
            try:
                os.unlink(temp_credentials_path)
            except OSError:
                pass
        for key, value in saved_env.items():
            if value is None:
                os.environ.pop(key, None)
            else:
                os.environ[key] = value


@pytest.fixture(scope="session")
def event_loop():
    try:
        loop = asyncio.get_running_loop()
    except RuntimeError:
        loop = asyncio.new_event_loop()
    yield loop
    loop.close()


@pytest.fixture(scope="function", autouse=True)
def setup_and_teardown(request):
    """
    This fixture reloads litellm before every function. To speed up testing by removing callbacks being chained.
    """


    if "google_genai_proxy_url" not in request.fixturenames:
        importlib.reload(litellm)

    loop = asyncio.get_event_loop_policy().new_event_loop()
    asyncio.set_event_loop(loop)
    print(litellm)
    yield

    # Teardown code (executes after the yield point)
    loop.close()  # Close the loop created earlier
    asyncio.set_event_loop(None)  # Remove the reference to the loop


@pytest.fixture(scope="module")
def vcr_config():
    return vcr_config_dict()


def pytest_recording_configure(config, vcr):
    register_persister_if_enabled(vcr)


@pytest.hookimpl(hookwrapper=True)
def pytest_runtest_makereport(item, call):
    outcome = yield
    rep = outcome.get_result()
    setattr(item, f"rep_{rep.when}", rep)


@pytest.fixture(autouse=True)
def _vcr_outcome_gate(request, vcr):
    install_live_call_probe(request, vcr)
    yield
    record_vcr_outcome(request, vcr)


def pytest_configure(config):
    _verbose_state.remember_pluginmanager(config)
    reset_vcr_diag_dir()


def pytest_runtest_logreport(report):
    _verbose_state.maybe_emit_verdict(report)


def pytest_collection_modifyitems(config, items):
    apply_vcr_auto_marker_to_items(
        items,
        skip_nodeid_suffixes=(
            "test_proxy_genai_sdk_non_streaming",
            "test_proxy_genai_sdk_streaming_completes_without_errors",
            "test_proxy_genai_sdk_streaming_dict_style",
        ),
    )

    # 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


def pytest_terminal_summary(terminalreporter, exitstatus, config):
    emit_cassette_cache_session_banner(terminalreporter)
    emit_vcr_classification_summary(terminalreporter)
    emit_vcr_diagnostic_log(terminalreporter)