litellm/tests/test_litellm/proxy/proxy_server/test_lifecycle.py

"""Behavior pins for proxy_server lifecycle, helpers, and small utilities.

Pins covered:
- ``proxy_startup_event``
- ``proxy_shutdown_event``
- ``_initialize_shared_aiohttp_session``
- ``cleanup_router_config_variables``
- ``save_worker_config``
- ``initialize``
- ``load_from_azure_key_vault``
- ``cost_tracking``
- ``_resolve_typed_dict_type``
- ``_resolve_pydantic_type``
- ``get_litellm_model_info``
- ``run_ollama_serve``
"""

from __future__ import annotations

import asyncio
import inspect
import json
import logging
import os
from collections.abc import Awaitable, Callable
from typing import List, Optional, Union
from unittest.mock import AsyncMock, MagicMock, patch

import pytest
from fastapi import FastAPI
from pydantic import BaseModel
from typing_extensions import TypedDict

import litellm.proxy.proxy_server as ps
from litellm.proxy.proxy_server import (
    ProxyStartupEvent,
    _initialize_shared_aiohttp_session,
    _resolve_pydantic_type,
    _resolve_typed_dict_type,
    cleanup_router_config_variables,
    cost_tracking,
    get_litellm_model_info,
    initialize,
    load_from_azure_key_vault,
    proxy_shutdown_event,
    proxy_startup_event,
    run_ollama_serve,
    save_worker_config,
)

from .conftest import normalize

# ---------------------------------------------------------------------------
# cleanup_router_config_variables
# ---------------------------------------------------------------------------


def test_cleanup_router_config_variables_resets_globals(monkeypatch):
    monkeypatch.setattr(ps, "master_key", "sk-sentinel", raising=False)
    monkeypatch.setattr(ps, "user_config_file_path", "/tmp/config.yaml", raising=False)
    monkeypatch.setattr(ps, "user_custom_auth", lambda x: x, raising=False)
    monkeypatch.setattr(ps, "health_check_interval", 42, raising=False)
    monkeypatch.setattr(ps, "prisma_client", MagicMock(), raising=False)

    cleanup_router_config_variables()

    observed = {
        "master_key": ps.master_key,
        "user_config_file_path": ps.user_config_file_path,
        "user_custom_auth": ps.user_custom_auth,
        "health_check_interval": ps.health_check_interval,
        "prisma_client": ps.prisma_client,
    }
    assert normalize(observed) == {
        "master_key": None,
        "user_config_file_path": None,
        "user_custom_auth": None,
        "health_check_interval": None,
        "prisma_client": None,
    }


def test_cleanup_router_config_variables_fails_on_unknown_attr_raises():
    """The function only writes documented globals — accessing a non-existent
    one after cleanup should still raise AttributeError."""
    cleanup_router_config_variables()
    with pytest.raises(AttributeError):
        _ = ps.this_attribute_should_not_exist_xyz


# ---------------------------------------------------------------------------
# proxy_shutdown_event
# ---------------------------------------------------------------------------


@pytest.mark.asyncio
async def test_proxy_shutdown_event_disconnects_prisma_and_resets(monkeypatch):
    fake_prisma = MagicMock()
    fake_prisma.disconnect = AsyncMock()
    monkeypatch.setattr(ps, "prisma_client", fake_prisma, raising=False)
    monkeypatch.setattr(ps, "master_key", "sk-x", raising=False)

    fake_jwt = MagicMock()
    fake_jwt.close = AsyncMock()
    monkeypatch.setattr(ps, "jwt_handler", fake_jwt, raising=False)
    monkeypatch.setattr(ps, "db_writer_client", None, raising=False)

    import litellm

    monkeypatch.setattr(litellm, "cache", None, raising=False)
    monkeypatch.setattr(litellm, "success_callback", [], raising=False)

    await proxy_shutdown_event()

    observed = {
        "disconnect_called": fake_prisma.disconnect.await_count == 1,
        "jwt_closed": fake_jwt.close.await_count == 1,
        "master_key_reset": ps.master_key,
        "prisma_reset": ps.prisma_client,
    }
    assert normalize(observed) == {
        "disconnect_called": True,
        "jwt_closed": True,
        "master_key_reset": None,
        "prisma_reset": None,
    }


@pytest.mark.asyncio
async def test_proxy_shutdown_drains_gateway_requests_before_disconnecting(monkeypatch):
    """
    The gateway request fold lives in memory, so shutdown drains it to the database.

    That drain has to happen while prisma is still connected: a write attempted
    after ``disconnect()`` raises ClientNotConnectedError, the flush swallows it
    and merges the counts back onto an accumulator the process is about to
    discard, and the final interval is lost silently on every restart. Ordering is
    the whole behavior here, so assert the order rather than that both ran.
    """
    calls: list = []  # mutable-ok: records call order, which is the assertion

    fake_prisma = MagicMock()
    fake_prisma.disconnect = AsyncMock(side_effect=lambda: calls.append("disconnect"))
    monkeypatch.setattr(ps, "prisma_client", fake_prisma, raising=False)

    async def _record_flush(client, accumulator):
        calls.append("flush")
        assert client is fake_prisma

    monkeypatch.setattr(ps, "flush_gateway_requests", _record_flush, raising=False)

    fake_jwt = MagicMock()
    fake_jwt.close = AsyncMock()
    monkeypatch.setattr(ps, "jwt_handler", fake_jwt, raising=False)
    monkeypatch.setattr(ps, "db_writer_client", None, raising=False)

    import litellm

    monkeypatch.setattr(litellm, "cache", None, raising=False)
    monkeypatch.setattr(litellm, "success_callback", [], raising=False)

    await proxy_shutdown_event()

    assert calls == ["flush", "disconnect"]


@pytest.mark.asyncio
async def test_proxy_shutdown_skips_gateway_flush_without_a_database(monkeypatch):
    """No prisma client means nothing to drain to, and no attempt is made."""
    flush = AsyncMock()
    monkeypatch.setattr(ps, "flush_gateway_requests", flush, raising=False)
    monkeypatch.setattr(ps, "prisma_client", None, raising=False)

    fake_jwt = MagicMock()
    fake_jwt.close = AsyncMock()
    monkeypatch.setattr(ps, "jwt_handler", fake_jwt, raising=False)
    monkeypatch.setattr(ps, "db_writer_client", None, raising=False)

    import litellm

    monkeypatch.setattr(litellm, "cache", None, raising=False)
    monkeypatch.setattr(litellm, "success_callback", [], raising=False)

    await proxy_shutdown_event()

    assert flush.await_count == 0


@pytest.mark.asyncio
async def test_proxy_shutdown_event_prisma_disconnect_raises_error(monkeypatch):
    fake_prisma = MagicMock()
    fake_prisma.disconnect = AsyncMock(side_effect=RuntimeError("db gone"))
    monkeypatch.setattr(ps, "prisma_client", fake_prisma, raising=False)

    fake_jwt = MagicMock()
    fake_jwt.close = AsyncMock()
    monkeypatch.setattr(ps, "jwt_handler", fake_jwt, raising=False)

    import litellm

    monkeypatch.setattr(litellm, "cache", None, raising=False)
    monkeypatch.setattr(litellm, "success_callback", [], raising=False)

    with pytest.raises(RuntimeError, match="db gone"):
        await proxy_shutdown_event()


# ---------------------------------------------------------------------------
# _flush_spend_logs_queue_on_shutdown
# ---------------------------------------------------------------------------


@pytest.mark.asyncio
async def test_flush_spend_logs_queue_on_shutdown_drains_before_disconnect(monkeypatch):
    fake_prisma = MagicMock()
    monkeypatch.setattr(ps, "prisma_client", fake_prisma, raising=False)
    monkeypatch.setattr(ps, "db_writer_client", None, raising=False)

    drain = AsyncMock()
    import litellm.proxy.utils as utils_mod

    monkeypatch.setattr(utils_mod, "drain_spend_logs_queue", drain)

    await ps._flush_spend_logs_queue_on_shutdown()

    observed = {
        "drain_calls": drain.await_count,
        "drain_prisma": drain.await_args.kwargs["prisma_client"] is fake_prisma,
    }
    assert observed == {
        "drain_calls": 1,
        "drain_prisma": True,
    }


@pytest.mark.asyncio
async def test_flush_spend_logs_queue_on_shutdown_swallows_drain_errors(monkeypatch):
    monkeypatch.setattr(ps, "prisma_client", MagicMock(), raising=False)
    monkeypatch.setattr(ps, "db_writer_client", None, raising=False)

    import litellm.proxy.utils as utils_mod

    monkeypatch.setattr(
        utils_mod,
        "drain_spend_logs_queue",
        AsyncMock(side_effect=RuntimeError("db gone")),
    )

    await ps._flush_spend_logs_queue_on_shutdown()


# ---------------------------------------------------------------------------
# _initialize_shared_aiohttp_session
# ---------------------------------------------------------------------------


@pytest.mark.asyncio
async def test_initialize_shared_aiohttp_session_returns_client_session():
    from aiohttp import ClientSession

    session = await _initialize_shared_aiohttp_session()
    try:
        observed = {
            "is_client_session": isinstance(session, ClientSession),
            "is_closed": session.closed,
            "has_connector": session.connector is not None,
        }
        assert normalize(observed) == {
            "is_client_session": True,
            "is_closed": False,
            "has_connector": True,
        }
    finally:
        if session is not None:
            await session.close()


@pytest.mark.asyncio
async def test_initialize_shared_aiohttp_session_aiohttp_missing_returns_none_on_failure(
    monkeypatch,
):
    """If aiohttp import fails, the function catches and returns None — no raise."""
    import builtins

    real_import = builtins.__import__

    def _raise_for_aiohttp(name, *args, **kwargs):
        if name == "aiohttp":
            raise ImportError("simulated missing aiohttp")
        return real_import(name, *args, **kwargs)

    monkeypatch.setattr(builtins, "__import__", _raise_for_aiohttp)
    result = await _initialize_shared_aiohttp_session()
    assert result is None


# ---------------------------------------------------------------------------
# save_worker_config
# ---------------------------------------------------------------------------


def test_save_worker_config_writes_json_to_environ(monkeypatch):
    monkeypatch.delenv("WORKER_CONFIG", raising=False)

    save_worker_config(model="gpt-4", config="/tmp/c.yaml", debug=True)

    payload = json.loads(os.environ["WORKER_CONFIG"])
    assert normalize(payload) == {
        "model": "gpt-4",
        "config": "/tmp/c.yaml",
        "debug": True,
    }


def test_save_worker_config_invalid_no_kwargs_yields_empty(monkeypatch):
    monkeypatch.delenv("WORKER_CONFIG", raising=False)

    save_worker_config()
    assert os.environ["WORKER_CONFIG"] == "{}"


# ---------------------------------------------------------------------------
# _redact_worker_config_for_logging (LIT-4152)
# ---------------------------------------------------------------------------


_LIT4152_SECRETS = (
    "sk-lit4152-regression-master-key-abcdef1234567890",
    "leak_password_9090",
    "sk-lit4152-provider-api-key-abcdef",
    "postgresql://leak_user:leak_password_9090@leak-host.internal:5432/leak_db",
)


def _lit4152_worker_config_dict():
    return {
        "model": "openai/gpt-4o-mini",
        "config": "/tmp/c.yaml",
        "master_key": _LIT4152_SECRETS[0],
        "database_url": _LIT4152_SECRETS[3],
        "api_key": _LIT4152_SECRETS[2],
        "telemetry": True,
    }


def test__redact_worker_config_for_logging_dict_masks_all_secret_shapes():
    """LIT-4152 regression: dict-form worker_config must not embed any raw
    secret. Covers the segment-matched fields (`master_key`, `api_key`) and the
    URL-with-credentials field (`database_url`), which the segment masker
    misses because neither segment matches its sensitive-pattern set.
    """
    from litellm.proxy.proxy_server import _redact_worker_config_for_logging

    redacted = _redact_worker_config_for_logging(_lit4152_worker_config_dict())
    rendered = repr(redacted)
    for secret in _LIT4152_SECRETS:
        assert secret not in rendered, f"leak: {secret} in {rendered!r}"
    assert isinstance(redacted, dict)
    assert redacted["model"] == "openai/gpt-4o-mini"
    assert redacted["telemetry"] is True


def test__redact_worker_config_for_logging_json_string_round_trips_masked():
    """Docker/K8s deployments hand the proxy a JSON string via ``WORKER_CONFIG``.
    Confirm the string path also masks and that the returned value re-parses
    into a dict with the sensitive fields masked.
    """
    from litellm.proxy.proxy_server import _redact_worker_config_for_logging

    payload = json.dumps(_lit4152_worker_config_dict())
    redacted = _redact_worker_config_for_logging(payload)
    assert isinstance(redacted, str)
    for secret in _LIT4152_SECRETS:
        assert secret not in redacted, f"leak: {secret} in {redacted!r}"
    parsed = json.loads(redacted)
    assert parsed["model"] == "openai/gpt-4o-mini"


def test__redact_worker_config_for_logging_passthrough_for_none_and_non_json_string():
    """Non-dict, non-JSON-parseable string is passed through verbatim (nothing
    to mask) and ``None`` returns ``None``.
    """
    from litellm.proxy.proxy_server import _redact_worker_config_for_logging

    assert _redact_worker_config_for_logging(None) is None
    assert _redact_worker_config_for_logging("/tmp/some_config.yaml") == "/tmp/some_config.yaml"


def test__redact_worker_config_for_logging_masks_non_string_url_webhook_values():
    """The URL/webhook fields the segment masker cannot catch by key name
    (``alert_to_webhook_url``, ``pass_through_endpoints``,
    ``database_extra_connection_params``) can hold non-string shapes:
    ``alert_to_webhook_url`` is typed as ``Optional[Dict]`` and can nest
    secret query params under keys the segment masker also misses. Confirm
    the whole value is replaced regardless of shape so a nested webhook or
    Bearer token under a non-segment-matched key does not slip through.
    """
    from litellm.proxy.proxy_server import _redact_worker_config_for_logging

    nested_webhook_secret = "https://hooks.slack.com/services/T0/B0/nested-webhook-secret-xyz"
    data = {
        "master_key": "sk-should-be-masked",
        "alert_to_webhook_url": {"budget_alerts": nested_webhook_secret},
        "pass_through_endpoints": [
            {
                "path": "/upstream",
                "target": "https://api.provider.com",
                "headers": {"Authorization": "Bearer nested-token-should-be-gone"},
            }
        ],
        "database_extra_connection_params": {"password": "extra-db-password-abc"},
    }
    redacted = _redact_worker_config_for_logging(data)
    rendered = repr(redacted)
    for secret in (
        "sk-should-be-masked",
        nested_webhook_secret,
        "nested-token-should-be-gone",
        "extra-db-password-abc",
    ):
        assert secret not in rendered, f"leak: {secret} in {rendered!r}"


def test__redact_worker_config_for_logging_masks_nested_secret_fields():
    """LIT-4152 nested regression: the URL/webhook credential fields the segment
    masker cannot catch by name (``database_url``,
    ``database_extra_connection_params``, ``pass_through_endpoints``,
    ``alert_to_webhook_url``) must be redacted at any depth, not just the top
    level. A worker_config that nests ``general_settings`` under a parent key
    must not leak a nested ``database_url`` or webhook secret; the earlier
    top-level-only redaction would have passed these through raw.
    """
    from litellm.proxy.proxy_server import _redact_worker_config_for_logging

    nested_db_url = "postgresql://nested_user:nested_pw_4152@nested-host:5432/db"
    nested_webhook = "https://hooks.slack.com/services/T0/B0/nested-4152-webhook"
    nested_extra_pw = "nested-extra-conn-pw-4152"
    nested_bearer = "Bearer nested-passthrough-token-4152"
    data = {
        "config": {
            "general_settings": {
                "database_url": nested_db_url,
                "database_extra_connection_params": {"password": nested_extra_pw},
                "alert_to_webhook_url": {"budget_alerts": nested_webhook},
                "pass_through_endpoints": [{"path": "/up", "headers": {"Authorization": nested_bearer}}],
            }
        }
    }
    redacted = _redact_worker_config_for_logging(data)
    rendered = repr(redacted)
    for secret in (nested_db_url, nested_webhook, nested_extra_pw, nested_bearer):
        assert secret not in rendered, f"nested leak: {secret} in {rendered!r}"

    inner = redacted["config"]["general_settings"]
    assert inner["database_url"] == "REDACTED"
    assert inner["pass_through_endpoints"] == "REDACTED"


# ---------------------------------------------------------------------------
# initialize
# ---------------------------------------------------------------------------


def test_initialize_signature_is_async_with_expected_params():
    sig = inspect.signature(initialize)
    # Hard-coded so a signature change (param added/removed) trips the gate.
    expected_param_count = 17
    observed = {
        "is_async": inspect.iscoroutinefunction(initialize),
        "param_count": len(sig.parameters),
        "has_model": "model" in sig.parameters,
        "has_config": "config" in sig.parameters,
    }
    assert normalize(observed) == {
        "is_async": True,
        "param_count": expected_param_count,
        "has_model": True,
        "has_config": True,
    }


@pytest.mark.asyncio
async def test_initialize_invalid_unexpected_kwarg_raises_type_error():
    with pytest.raises(TypeError):
        await initialize(this_is_not_a_real_kwarg=True)


# ---------------------------------------------------------------------------
# load_from_azure_key_vault
# ---------------------------------------------------------------------------


def test_load_from_azure_key_vault_disabled_no_side_effect(monkeypatch):
    import litellm

    sentinel_secret_mgr = object()
    monkeypatch.setattr(litellm, "secret_manager_client", sentinel_secret_mgr, raising=False)

    result = load_from_azure_key_vault(use_azure_key_vault=False)

    observed = {
        "return_value": result,
        "secret_manager_unchanged": litellm.secret_manager_client is sentinel_secret_mgr,
        "called_with": False,
    }
    assert normalize(observed) == {
        "return_value": None,
        "secret_manager_unchanged": True,
        "called_with": False,
    }


def test_load_from_azure_key_vault_missing_uri_failure_is_swallowed(monkeypatch):
    """Enabled but AZURE_KEY_VAULT_URI unset / azure libs likely unavailable —
    function catches Exception and does not raise."""
    monkeypatch.delenv("AZURE_KEY_VAULT_URI", raising=False)

    result = load_from_azure_key_vault(use_azure_key_vault=True)
    assert result is None


# ---------------------------------------------------------------------------
# cost_tracking
# ---------------------------------------------------------------------------


def test_cost_tracking_adds_db_and_shadow_eval_callbacks_when_prisma_set(monkeypatch):
    import litellm
    from litellm.integrations.shadow_eval_logger import ShadowEvalLogger

    fake_prisma = MagicMock()
    monkeypatch.setattr(ps, "prisma_client", fake_prisma, raising=False)
    monkeypatch.setattr(litellm, "callbacks", [], raising=False)
    monkeypatch.setattr(litellm, "_async_success_callback", [], raising=False)

    before_callbacks = len(litellm.callbacks)
    before_async = len(litellm._async_success_callback)

    cost_tracking()
    cost_tracking()

    observed = {
        "added_to_callbacks": len(litellm.callbacks) - before_callbacks,
        "added_to_async_success": len(litellm._async_success_callback) - before_async,
        "shadow_eval_loggers": sum(isinstance(cb, ShadowEvalLogger) for cb in litellm.callbacks),
        "prisma_was_set": True,
    }
    assert normalize(observed) == {
        "added_to_callbacks": 2,
        "added_to_async_success": 1,
        "shadow_eval_loggers": 1,
        "prisma_was_set": True,
    }


def test_cost_tracking_no_op_when_prisma_missing(monkeypatch):
    """Without a prisma_client cost_tracking is a no-op — not an error."""
    import litellm

    monkeypatch.setattr(ps, "prisma_client", None, raising=False)
    monkeypatch.setattr(litellm, "callbacks", [], raising=False)
    monkeypatch.setattr(litellm, "_async_success_callback", [], raising=False)

    cost_tracking()

    assert litellm.callbacks == []
    assert litellm._async_success_callback == []


# ---------------------------------------------------------------------------
# _resolve_typed_dict_type
# ---------------------------------------------------------------------------


class _SampleTD(TypedDict):
    a: int
    b: str


def test_resolve_typed_dict_type_finds_class_in_optional():
    typ = Optional[_SampleTD]
    result = _resolve_typed_dict_type(typ)

    observed = {
        "input_repr": "Optional[_SampleTD]",
        "result_is_sample_td": result is _SampleTD,
        "result_is_class": isinstance(result, type),
    }
    assert normalize(observed) == {
        "input_repr": "Optional[_SampleTD]",
        "result_is_sample_td": True,
        "result_is_class": True,
    }


def test_resolve_typed_dict_type_invalid_plain_type_returns_none():
    """A non-TypedDict, non-Union input returns None — not an error."""
    assert _resolve_typed_dict_type(int) is None
    assert _resolve_typed_dict_type(str) is None


# ---------------------------------------------------------------------------
# _resolve_pydantic_type
# ---------------------------------------------------------------------------


class _SampleModelA(BaseModel):
    x: int


class _SampleModelB(BaseModel):
    y: str


def test_resolve_pydantic_type_extracts_non_none_args_from_union():
    typ = Union[_SampleModelA, _SampleModelB, None]
    result = _resolve_pydantic_type(typ)

    observed = {
        "result_type": type(result).__name__,
        "result_len": len(result),
        "contains_a": _SampleModelA in result,
        "contains_b": _SampleModelB in result,
    }
    assert normalize(observed) == {
        "result_type": "list",
        "result_len": 2,
        "contains_a": True,
        "contains_b": True,
    }


def test_resolve_pydantic_type_invalid_non_union_non_model_returns_empty():
    """When given a non-Union and non-BaseModel input the function returns [].

    This is the silent-empty fallback path — error-ish by behavior."""
    result = _resolve_pydantic_type(int)
    assert result == []


# ---------------------------------------------------------------------------
# get_litellm_model_info
# ---------------------------------------------------------------------------


def test_get_litellm_model_info_uses_base_model_for_lookup(monkeypatch):
    import litellm

    expected_info = {"max_tokens": 8192, "input_cost_per_token": 0.00003}
    fake_get = MagicMock(return_value=expected_info)
    monkeypatch.setattr(litellm, "get_model_info", fake_get, raising=False)

    model = {
        "model_info": {"base_model": "gpt-4"},
        "litellm_params": {"model": "azure/my-deployment"},
    }
    result = get_litellm_model_info(model=model)

    observed = {
        "called_arg": (
            fake_get.call_args.args[0] if fake_get.call_args.args else fake_get.call_args.kwargs.get("model")
        ),
        "returned_max_tokens": result.get("max_tokens"),
        "returned_cost": result.get("input_cost_per_token"),
    }
    assert normalize(observed) == {
        "called_arg": "gpt-4",
        "returned_max_tokens": 8192,
        "returned_cost": 0.00003,
    }


def test_get_litellm_model_info_invalid_empty_dict_returns_empty():
    """Empty input means model_to_lookup is None — internal exception is caught
    and the function returns {}."""
    result = get_litellm_model_info(model={})
    assert result == {}


# ---------------------------------------------------------------------------
# run_ollama_serve
# ---------------------------------------------------------------------------


def test_run_ollama_serve_invokes_subprocess_popen(monkeypatch):
    fake_popen = MagicMock()
    monkeypatch.setattr(ps.subprocess, "Popen", fake_popen)

    run_ollama_serve()

    args, kwargs = fake_popen.call_args
    observed = {
        "popen_called": fake_popen.call_count == 1,
        "command": args[0] if args else kwargs.get("args"),
        "has_stdout_kw": "stdout" in kwargs,
        "has_stderr_kw": "stderr" in kwargs,
    }
    assert normalize(observed) == {
        "popen_called": True,
        "command": ["ollama", "serve"],
        "has_stdout_kw": True,
        "has_stderr_kw": True,
    }


def test_run_ollama_serve_popen_failure_is_swallowed(monkeypatch):
    """Popen raising OSError must NOT propagate — function logs and returns."""
    monkeypatch.setattr(ps.subprocess, "Popen", MagicMock(side_effect=OSError("no ollama binary")))

    result = run_ollama_serve()
    assert result is None


# ---------------------------------------------------------------------------
# proxy_startup_event
# ---------------------------------------------------------------------------


@pytest.mark.asyncio
async def test_proxy_startup_event_is_async_context_manager_with_expected_signature():
    """proxy_startup_event is the FastAPI lifespan. Verify its surface without
    actually running the heavy init path (DB, Router, OTEL, etc.)."""
    sig = inspect.signature(proxy_startup_event)
    wrapped = getattr(proxy_startup_event, "__wrapped__", None)
    observed = {
        "param_count": len(sig.parameters),
        "has_app_param": "app" in sig.parameters,
        "wrapped_is_async": inspect.iscoroutinefunction(wrapped) or inspect.isasyncgenfunction(wrapped),
        "has_asynccontextmanager_wrapper": wrapped is not None,
    }
    assert normalize(observed) == {
        "param_count": 1,
        "has_app_param": True,
        "wrapped_is_async": True,
        "has_asynccontextmanager_wrapper": True,
    }


@pytest.mark.asyncio
async def test_proxy_startup_event_invalid_missing_app_arg_raises():
    """Calling the lifespan with no FastAPI app argument must fail."""
    with pytest.raises(TypeError):
        # Intentionally invoke the underlying async generator function with
        # no arguments — the decorator preserves the missing-arg TypeError.
        async with proxy_startup_event():  # type: ignore[call-arg]
            pass


def test_otel_global_provider_published_after_callback_init():
    """The OTel V2 global-provider publish must run after callback
    initialization in ``proxy_startup_event``.

    Regression for the orphan span: a preset (arize, langfuse, …) builds its
    single folded logger during ``_initialize_startup_logging``. Publishing the
    global ``TracerProvider`` before that ran found no logger and built a second
    generic one whose provider became the global, so the FastAPI server span and
    the preset's gen-ai spans exported through different providers and the LLM
    span was orphaned. The publish (``publish_global_otel_v2_provider``) must
    therefore appear after ``_initialize_startup_logging`` in the lifespan source.
    """
    wrapped = getattr(proxy_startup_event, "__wrapped__", proxy_startup_event)
    source = inspect.getsource(wrapped)
    init_pos = source.find("_initialize_startup_logging(")
    publish_pos = source.find("publish_global_otel_v2_provider(")
    assert init_pos != -1, "callback init call not found in proxy_startup_event"
    assert publish_pos != -1, "OTEL global publish not found in proxy_startup_event"
    assert init_pos < publish_pos, (
        "OTEL global provider is published before callbacks are initialized; a "
        "preset logger will not exist yet and a second generic logger will own "
        "the global provider, orphaning gen-ai spans"
    )


def test_startup_warns_for_global_budget_without_database(caplog):
    with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"):
        ProxyStartupEvent._warn_budget_without_db(max_budget=100.0, prisma_client=None)

    assert "litellm.max_budget=100.0" in caplog.text
    assert "will NOT be enforced" in caplog.text
    assert "requests will never be blocked" in caplog.text


def test_startup_does_not_warn_for_global_budget_with_database(caplog):
    with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"):
        ProxyStartupEvent._warn_budget_without_db(max_budget=100.0, prisma_client=MagicMock())

    assert "litellm.max_budget" not in caplog.text


@pytest.mark.parametrize("max_budget", [0, None])
def test_startup_does_not_warn_without_global_budget(caplog, max_budget):
    with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"):
        ProxyStartupEvent._warn_budget_without_db(max_budget=max_budget, prisma_client=None)

    assert "litellm.max_budget" not in caplog.text


def test_proxy_startup_event_warns_for_global_budget_without_database():
    """Pin the lifespan call that prevents silent DB-less budgets.

    The call must follow Prisma setup so DB-backed deployments do not false-positive.
    Direct ``_warn_budget_without_db`` tests cover the warning behavior itself.
    """
    wrapped = getattr(proxy_startup_event, "__wrapped__", proxy_startup_event)
    source = inspect.getsource(wrapped)
    budget_check_pos = source.find("if prisma_client is not None and litellm.max_budget > 0:")
    warn_pos = source.find("_warn_budget_without_db(")
    next_startup_section_pos = source.find(
        "await ProxyStartupEvent.initialize_scheduled_background_jobs(",
        budget_check_pos,
    )

    assert budget_check_pos != -1, "global budget startup block not found"
    assert warn_pos != -1, "DB-less budget warning call not found"
    assert next_startup_section_pos != -1, "startup section after budget block not found"
    assert budget_check_pos < warn_pos < next_startup_section_pos, (
        "DB-less budget warning must run after Prisma setup and the DB-backed budget block"
    )


# ---------------------------------------------------------------------------
# _initialize_slack_alerting_jobs — spend-report pod locking (issue #14809)
# ---------------------------------------------------------------------------

SlackAlertingJobs = dict[str, Callable[[], Awaitable[None]]]


def _make_slack_alerting_proxy_logging(acquire_lock_result: bool | None) -> MagicMock:
    proxy_logging_obj = MagicMock()
    proxy_logging_obj.slack_alerting_instance.alerting = ["slack"]
    proxy_logging_obj.slack_alerting_instance.send_weekly_spend_report = AsyncMock()
    proxy_logging_obj.slack_alerting_instance.send_monthly_spend_report = AsyncMock()
    proxy_logging_obj.slack_alerting_instance.send_fallback_stats_from_prometheus = AsyncMock()
    pod_lock_manager = proxy_logging_obj.db_spend_update_writer.pod_lock_manager
    pod_lock_manager.acquire_lock = AsyncMock(return_value=acquire_lock_result)
    pod_lock_manager.release_lock = AsyncMock()
    return proxy_logging_obj


async def _init_slack_alerting_jobs(
    acquire_lock_result: bool | None,
    spend_report_frequency: str = "7d",
) -> tuple[SlackAlertingJobs, MagicMock]:
    scheduler = MagicMock()
    proxy_logging_obj = _make_slack_alerting_proxy_logging(acquire_lock_result)

    await ProxyStartupEvent._initialize_slack_alerting_jobs(
        scheduler=scheduler,
        general_settings={"spend_report_frequency": spend_report_frequency},
        proxy_logging_obj=proxy_logging_obj,
        prisma_client=MagicMock(),
    )

    jobs = {call.kwargs["id"]: call.args[0] for call in scheduler.add_job.call_args_list}
    return jobs, proxy_logging_obj


@pytest.mark.parametrize("spend_report_frequency", ["0d", "-1d", "7h"])
@pytest.mark.asyncio
async def test_initialize_slack_alerting_jobs_invalid_frequency_raises(spend_report_frequency: str):
    """A non-positive window used to become an every-second APScheduler interval, and now also
    computes a negative lock TTL that expires instantly and suppresses the report for good.
    match= is load-bearing: drop the guard and "-1d" still raises, but from duration_in_seconds."""
    with pytest.raises(ValueError, match="positive number of days"):
        await _init_slack_alerting_jobs(
            acquire_lock_result=True,
            spend_report_frequency=spend_report_frequency,
        )


@pytest.mark.asyncio
async def test_weekly_spend_report_skipped_when_another_pod_holds_the_lock():
    """regression: issue #14809 - every pod ran its own weekly spend report job."""
    jobs, proxy_logging_obj = await _init_slack_alerting_jobs(acquire_lock_result=False)

    await jobs["weekly_spend_report_job"]()

    proxy_logging_obj.slack_alerting_instance.send_weekly_spend_report.assert_not_awaited()
    proxy_logging_obj.db_spend_update_writer.pod_lock_manager.acquire_lock.assert_awaited_once_with(
        cronjob_id="weekly_spend_report_job",
        ttl=7 * 86400 - 3600,
        allow_reentrant=False,
    )


@pytest.mark.parametrize("acquire_lock_result", [True, None])
@pytest.mark.asyncio
async def test_weekly_spend_report_sent_when_the_lock_is_free_or_absent(acquire_lock_result):
    """None means redis isn't configured; a single-pod deploy must still report."""
    jobs, proxy_logging_obj = await _init_slack_alerting_jobs(acquire_lock_result=acquire_lock_result)

    await jobs["weekly_spend_report_job"]()

    proxy_logging_obj.slack_alerting_instance.send_weekly_spend_report.assert_awaited_once_with("7d")


@pytest.mark.asyncio
async def test_weekly_spend_report_lock_ttl_tracks_the_configured_window():
    """TTL is the window less an hour: long enough that no second pod re-sends inside the
    window, short enough that the lock is gone before the next one opens. A fixed TTL would
    break one end or the other as soon as spend_report_frequency changes."""
    jobs, proxy_logging_obj = await _init_slack_alerting_jobs(acquire_lock_result=True, spend_report_frequency="1d")

    await jobs["weekly_spend_report_job"]()

    proxy_logging_obj.db_spend_update_writer.pod_lock_manager.acquire_lock.assert_awaited_once_with(
        cronjob_id="weekly_spend_report_job",
        ttl=86400 - 3600,
        allow_reentrant=False,
    )
    proxy_logging_obj.slack_alerting_instance.send_weekly_spend_report.assert_awaited_once_with("1d")


@pytest.mark.asyncio
async def test_monthly_spend_report_skipped_when_another_pod_holds_the_lock():
    jobs, proxy_logging_obj = await _init_slack_alerting_jobs(acquire_lock_result=False)

    await jobs["monthly_spend_report_job"]()

    proxy_logging_obj.slack_alerting_instance.send_monthly_spend_report.assert_not_awaited()
    proxy_logging_obj.db_spend_update_writer.pod_lock_manager.acquire_lock.assert_awaited_once_with(
        cronjob_id="monthly_spend_report_job",
        ttl=3600,
        allow_reentrant=False,
    )


@pytest.mark.parametrize("acquire_lock_result", [True, None])
@pytest.mark.asyncio
async def test_monthly_spend_report_sent_when_the_lock_is_free_or_absent(acquire_lock_result):
    jobs, proxy_logging_obj = await _init_slack_alerting_jobs(acquire_lock_result=acquire_lock_result)

    await jobs["monthly_spend_report_job"]()

    proxy_logging_obj.slack_alerting_instance.send_monthly_spend_report.assert_awaited_once_with()


@pytest.mark.asyncio
async def test_spend_report_locks_are_never_released():
    """The lock is a per-window marker, not a mutex: releasing it lets the next pod re-send."""
    jobs, proxy_logging_obj = await _init_slack_alerting_jobs(acquire_lock_result=True)

    await jobs["weekly_spend_report_job"]()
    await jobs["monthly_spend_report_job"]()

    proxy_logging_obj.db_spend_update_writer.pod_lock_manager.release_lock.assert_not_awaited()


@pytest.mark.asyncio
async def test_prometheus_fallback_stats_job_skipped_when_another_pod_holds_the_lock(monkeypatch):
    """The boot-time send goes through the same gate, so a losing pod sends nothing at all:
    startup and the scheduled job both stay at zero."""
    monkeypatch.setenv("PROMETHEUS_URL", "http://prometheus.invalid")
    jobs, proxy_logging_obj = await _init_slack_alerting_jobs(acquire_lock_result=False)
    send_fallback_stats = proxy_logging_obj.slack_alerting_instance.send_fallback_stats_from_prometheus
    assert send_fallback_stats.await_count == 0

    await jobs["prometheus_fallback_stats_job"]()

    assert send_fallback_stats.await_count == 0
    proxy_logging_obj.db_spend_update_writer.pod_lock_manager.acquire_lock.assert_awaited_with(
        cronjob_id="prometheus_fallback_stats_job",
        ttl=3600,
        allow_reentrant=False,
    )
    assert proxy_logging_obj.db_spend_update_writer.pod_lock_manager.acquire_lock.await_count == 2


@pytest.mark.parametrize("acquire_lock_result", [True, None])
@pytest.mark.asyncio
async def test_prometheus_fallback_stats_job_runs_when_the_lock_is_free_or_absent(monkeypatch, acquire_lock_result):
    monkeypatch.setenv("PROMETHEUS_URL", "http://prometheus.invalid")
    jobs, proxy_logging_obj = await _init_slack_alerting_jobs(acquire_lock_result=acquire_lock_result)
    send_fallback_stats = proxy_logging_obj.slack_alerting_instance.send_fallback_stats_from_prometheus
    assert send_fallback_stats.await_count == 1

    await jobs["prometheus_fallback_stats_job"]()

    assert send_fallback_stats.await_count == 2