litellm/tests/test_litellm/proxy/utils/proxy_logging/test_lifecycle.py

"""Pin ProxyLogging lifecycle: ``__init__``, ``startup_event``,
``update_values``, ``_add_proxy_hooks``, ``get_proxy_hook``, and
``_init_litellm_callbacks``.

Also covers ``update_request_status`` and ``_convert_user_api_key_auth_to_dict``
because they are direct dependents on the lifecycle state.
"""

from __future__ import annotations

import asyncio
from typing import Any, Dict, List
from unittest.mock import AsyncMock, MagicMock, patch

import pytest

import litellm
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
from litellm.proxy.utils import (
    InternalUsageCache,
    ProxyLogging,
)


# ---------------------------------------------------------------------------
# __init__
# ---------------------------------------------------------------------------


def test_proxy_logging_init_sets_default_state(mock_callbacks_disabled):
    cache = UserApiKeyCache()
    pl = ProxyLogging(user_api_key_cache=cache)
    snapshot = {
        "internal_usage_cache_type": type(pl.internal_usage_cache).__name__,
        "alerting_is_none": pl.alerting is None,
        "alerting_threshold": pl.alerting_threshold,
        "premium_user": pl.premium_user,
        "proxy_hook_mapping": pl.proxy_hook_mapping,
        "daily_report_started": pl.daily_report_started,
        "hanging_requests_check_started": pl.hanging_requests_check_started,
    }
    assert snapshot == {
        "internal_usage_cache_type": "InternalUsageCache",
        "alerting_is_none": True,
        "alerting_threshold": 300,
        "premium_user": False,
        "proxy_hook_mapping": {},
        "daily_report_started": False,
        "hanging_requests_check_started": False,
    }


def test_proxy_logging_init_premium_user_flag(mock_callbacks_disabled):
    pl = ProxyLogging(user_api_key_cache=UserApiKeyCache(), premium_user=True)
    assert pl.premium_user is True


def test_proxy_logging_init_missing_cache_raises():
    with pytest.raises(TypeError):
        ProxyLogging()  # type: ignore[call-arg]


# ---------------------------------------------------------------------------
# update_values
# ---------------------------------------------------------------------------


def test_update_values_stores_alerting_state(proxy_logging):
    proxy_logging.slack_alerting_instance = MagicMock()
    proxy_logging.update_values(
        alerting=["slack"],
        alerting_threshold=42.0,
        alert_types=["llm_too_slow"],
        alert_to_webhook_url={"key": "value"},
    )
    snapshot = {
        "alerting": proxy_logging.alerting,
        "threshold": proxy_logging.alerting_threshold,
        "alert_types": proxy_logging.alert_types,
        "webhook_url": proxy_logging.alert_to_webhook_url,
    }
    assert snapshot == {
        "alerting": ["slack"],
        "threshold": 42.0,
        "alert_types": ["llm_too_slow"],
        "webhook_url": {"key": "value"},
    }


def test_update_values_with_only_redis_cache_does_not_touch_slack(proxy_logging):
    proxy_logging.slack_alerting_instance = MagicMock()
    redis = MagicMock()
    proxy_logging.update_values(redis_cache=redis)
    proxy_logging.slack_alerting_instance.update_values.assert_not_called()
    assert proxy_logging.internal_usage_cache.dual_cache.redis_cache is redis


def test_update_values_with_no_args_is_noop(proxy_logging):
    proxy_logging.slack_alerting_instance = MagicMock()
    proxy_logging.update_values()
    proxy_logging.slack_alerting_instance.update_values.assert_not_called()


def test_update_values_invalid_type_for_alerting_raises(proxy_logging):
    proxy_logging.slack_alerting_instance = MagicMock(
        update_values=MagicMock(side_effect=TypeError("bad type"))
    )
    with pytest.raises(TypeError):
        proxy_logging.update_values(alerting={"not": "a list"})  # type: ignore[arg-type]


# ---------------------------------------------------------------------------
# startup_event
# ---------------------------------------------------------------------------


def test_startup_event_initializes_slack_and_callbacks(proxy_logging):
    proxy_logging.slack_alerting_instance = MagicMock()
    proxy_logging.slack_alerting_instance.alert_types = []
    proxy_logging._init_litellm_callbacks = MagicMock()
    proxy_logging.update_values = MagicMock()

    proxy_logging.startup_event(llm_router=None, redis_usage_cache=None)
    snapshot = {
        "update_called": proxy_logging.update_values.called,
        "init_called": proxy_logging._init_litellm_callbacks.called,
        "slack_update_called": proxy_logging.slack_alerting_instance.update_values.called,
    }
    assert snapshot == {
        "update_called": True,
        "init_called": True,
        "slack_update_called": True,
    }


def test_startup_event_propagates_init_callbacks_failure_raises(proxy_logging):
    proxy_logging.slack_alerting_instance = MagicMock()
    proxy_logging.slack_alerting_instance.alert_types = []
    proxy_logging._init_litellm_callbacks = MagicMock(side_effect=RuntimeError("boom"))

    with pytest.raises(RuntimeError, match="boom"):
        proxy_logging.startup_event(llm_router=None, redis_usage_cache=None)


# ---------------------------------------------------------------------------
# _add_proxy_hooks
# ---------------------------------------------------------------------------


def test_add_proxy_hooks_registers_callbacks(proxy_logging, monkeypatch):
    """Patch ``PROXY_HOOKS`` and the resolver so we control exactly
    what gets registered. Verifies that the resulting instances land in
    ``proxy_logging.proxy_hook_mapping`` keyed by hook name.
    """
    hook_keys = ["cache_control_check", "max_budget_limiter"]
    registered: List[Any] = []

    from litellm.proxy import utils as utils_mod

    def fake_get_proxy_hook(hook_name):
        class _Stub:
            __name__ = hook_name

            def __init__(self, **kwargs):
                self.hook_name = hook_name

        return _Stub

    monkeypatch.setattr(utils_mod, "PROXY_HOOKS", hook_keys)
    monkeypatch.setattr(utils_mod, "get_proxy_hook", fake_get_proxy_hook)
    monkeypatch.setattr(
        litellm.logging_callback_manager,
        "add_litellm_callback",
        lambda cb: registered.append(cb),
    )

    with patch("litellm.proxy.proxy_server.prisma_client", None):
        proxy_logging._add_proxy_hooks(llm_router=None)

    keys = list(proxy_logging.proxy_hook_mapping.keys())
    snapshot = {
        "mapping_keys": keys,
        "registered_count": len(registered),
        "registered_hook_names": [getattr(r, "hook_name", None) for r in registered],
    }
    assert snapshot == {
        "mapping_keys": hook_keys,
        "registered_count": len(hook_keys),
        "registered_hook_names": hook_keys,
    }


def test_add_proxy_hooks_unknown_hook_raises(proxy_logging, monkeypatch):
    from litellm.proxy import utils as utils_mod

    monkeypatch.setattr(utils_mod, "PROXY_HOOKS", ["bogus_hook"])

    def bad_resolver(name):
        raise KeyError(name)

    monkeypatch.setattr(utils_mod, "get_proxy_hook", bad_resolver)
    with pytest.raises(KeyError):
        proxy_logging._add_proxy_hooks(llm_router=None)


# ---------------------------------------------------------------------------
# get_proxy_hook
# ---------------------------------------------------------------------------


def test_get_proxy_hook_returns_registered_instance(proxy_logging):
    s_cache = MagicMock()
    s_budget = MagicMock()
    s_parallel = MagicMock()
    proxy_logging.proxy_hook_mapping = {
        "cache_control_check": s_cache,
        "max_budget_limiter": s_budget,
        "max_parallel_request_limiter": s_parallel,
    }
    snapshot = {
        "cache_control_check": proxy_logging.get_proxy_hook("cache_control_check") is s_cache,
        "max_budget_limiter": proxy_logging.get_proxy_hook("max_budget_limiter") is s_budget,
        "max_parallel_request_limiter": proxy_logging.get_proxy_hook("max_parallel_request_limiter") is s_parallel,
        "unknown_returns_none": proxy_logging.get_proxy_hook("unknown") is None,
    }
    assert snapshot == {
        "cache_control_check": True,
        "max_budget_limiter": True,
        "max_parallel_request_limiter": True,
        "unknown_returns_none": True,
    }


def test_get_proxy_hook_unknown_returns_none(proxy_logging):
    proxy_logging.proxy_hook_mapping = {}
    assert proxy_logging.get_proxy_hook("does-not-exist") is None


def test_get_proxy_hook_non_string_key_raises(proxy_logging):
    # ``dict.get`` doesn't raise on unhashable types — but ``None`` returns None.
    # The pin: passing an unhashable key blows up like dict access does.
    proxy_logging.proxy_hook_mapping = {"k": object()}
    with pytest.raises(TypeError):
        proxy_logging.get_proxy_hook({"unhashable": True})  # type: ignore[arg-type]


# ---------------------------------------------------------------------------
# _init_litellm_callbacks
# ---------------------------------------------------------------------------


def test_init_litellm_callbacks_replaces_string_with_instance(proxy_logging, monkeypatch):
    from litellm.proxy import utils as utils_mod

    sentinel_instance = MagicMock(spec=litellm.integrations.custom_logger.CustomLogger)
    sentinel_instance.__class__ = litellm.integrations.custom_logger.CustomLogger

    monkeypatch.setattr(litellm, "callbacks", ["some-string-logger"])
    monkeypatch.setattr(
        litellm.litellm_core_utils.litellm_logging,
        "_init_custom_logger_compatible_class",
        lambda *a, **kw: sentinel_instance,
    )

    monkeypatch.setattr(utils_mod, "PROXY_HOOKS", [])
    proxy_logging._init_litellm_callbacks(llm_router=None)
    snapshot = {
        "replaced_first_item": litellm.callbacks[0] is sentinel_instance,
        "callbacks_grew_with_service": len(litellm.callbacks) >= 2,
        "service_logging_appended": any(
            "ServiceLogging" in type(c).__name__ for c in litellm.callbacks
        ),
    }
    assert snapshot == {
        "replaced_first_item": True,
        "callbacks_grew_with_service": True,
        "service_logging_appended": True,
    }


def test_init_litellm_callbacks_string_resolution_failure_keeps_string(proxy_logging, monkeypatch):
    from litellm.proxy import utils as utils_mod

    monkeypatch.setattr(litellm, "callbacks", ["unknown-logger"])
    monkeypatch.setattr(
        litellm.litellm_core_utils.litellm_logging,
        "_init_custom_logger_compatible_class",
        lambda *a, **kw: None,
    )
    monkeypatch.setattr(utils_mod, "PROXY_HOOKS", [])
    proxy_logging._init_litellm_callbacks(llm_router=None)
    # Resolver returned None — original string remains in place at idx 0.
    assert litellm.callbacks[0] == "unknown-logger"


def test_init_litellm_callbacks_propagates_resolver_error_raises(proxy_logging, monkeypatch):
    from litellm.proxy import utils as utils_mod

    monkeypatch.setattr(litellm, "callbacks", ["raises-on-init"])
    monkeypatch.setattr(
        litellm.litellm_core_utils.litellm_logging,
        "_init_custom_logger_compatible_class",
        MagicMock(side_effect=RuntimeError("bad init")),
    )
    monkeypatch.setattr(utils_mod, "PROXY_HOOKS", [])
    with pytest.raises(RuntimeError):
        proxy_logging._init_litellm_callbacks(llm_router=None)


# ---------------------------------------------------------------------------
# update_request_status
# ---------------------------------------------------------------------------


@pytest.mark.asyncio
async def test_update_request_status_when_alerting_set_writes_cache(proxy_logging):
    proxy_logging.alerting = ["slack"]
    proxy_logging.alerting_threshold = 5.0
    captured: Dict[str, Any] = {}

    async def fake_set_cache(**kwargs):
        captured.update(kwargs)

    proxy_logging.internal_usage_cache.async_set_cache = fake_set_cache  # type: ignore[assignment]
    await proxy_logging.update_request_status(litellm_call_id="call-1", status="success")
    snapshot = {
        "key": captured["key"],
        "value": captured["value"],
        "local_only": captured["local_only"],
        "ttl": captured["ttl"],
    }
    assert snapshot == {
        "key": "request_status:call-1",
        "value": "success",
        "local_only": True,
        "ttl": 105.0,
    }


@pytest.mark.asyncio
async def test_update_request_status_no_alerting_skips_cache(proxy_logging):
    proxy_logging.alerting = None
    proxy_logging.internal_usage_cache.async_set_cache = AsyncMock()
    await proxy_logging.update_request_status(litellm_call_id="call-1", status="success")
    proxy_logging.internal_usage_cache.async_set_cache.assert_not_called()


@pytest.mark.asyncio
async def test_update_request_status_cache_error_raises(proxy_logging):
    proxy_logging.alerting = ["slack"]
    proxy_logging.internal_usage_cache.async_set_cache = AsyncMock(side_effect=ConnectionError("redis"))
    with pytest.raises(ConnectionError):
        await proxy_logging.update_request_status(litellm_call_id="x", status="fail")


# ---------------------------------------------------------------------------
# _convert_user_api_key_auth_to_dict
# ---------------------------------------------------------------------------


def test_convert_user_api_key_auth_to_dict_pydantic_uses_model_dump(proxy_logging, make_user_api_key_auth):
    auth = make_user_api_key_auth(user_id="u-1", team_id="t-1")
    result = proxy_logging._convert_user_api_key_auth_to_dict(auth)
    snapshot = {
        "user_id": result["user_id"],
        "team_id": result["team_id"],
        "is_dict": isinstance(result, dict),
    }
    assert snapshot == {"user_id": "u-1", "team_id": "t-1", "is_dict": True}


def test_convert_user_api_key_auth_to_dict_plain_object_uses_dict(proxy_logging):
    class Obj:
        pass

    obj = Obj()
    obj.a = 1
    obj.b = 2
    obj.c = 3
    result = proxy_logging._convert_user_api_key_auth_to_dict(obj)
    assert result == {"a": 1, "b": 2, "c": 3}


def test_convert_user_api_key_auth_to_dict_none_returns_empty_dict(proxy_logging):
    assert proxy_logging._convert_user_api_key_auth_to_dict(None) == {}


def test_convert_user_api_key_auth_to_dict_unconvertible_object_returns_empty(proxy_logging):
    class NoDict:
        __slots__ = ()

    assert proxy_logging._convert_user_api_key_auth_to_dict(NoDict()) == {}


def test_convert_user_api_key_auth_to_dict_pydantic_error_raises(proxy_logging):
    """A ``model_dump`` that raises propagates."""

    class _Boom:
        def model_dump(self):
            raise RuntimeError("model_dump failure")

    with pytest.raises(RuntimeError):
        proxy_logging._convert_user_api_key_auth_to_dict(_Boom())