"""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 ( 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, } @pytest.mark.asyncio async def test_startup_event_schedules_deprecation_check_before_its_alert_type_is_on(proxy_logging): """Alerting config can enable the deprecation alert after startup, so the loop must already be running""" proxy_logging.alerting = ["slack"] proxy_logging.slack_alerting_instance = MagicMock() proxy_logging.slack_alerting_instance.alert_types = [] proxy_logging.slack_alerting_instance.run_scheduled_deprecation_check = AsyncMock() proxy_logging._init_litellm_callbacks = MagicMock() proxy_logging.startup_event(llm_router=None, redis_usage_cache=None) assert proxy_logging.deprecation_check_started is True proxy_logging.slack_alerting_instance.run_scheduled_deprecation_check.assert_called_once_with( pod_lock_manager=proxy_logging.db_spend_update_writer.pod_lock_manager ) @pytest.mark.asyncio async def test_update_values_schedules_deprecation_check_when_alerting_arrives_later(proxy_logging): """A proxy that boots without alerting still needs the loop once a config reload turns it on""" proxy_logging.slack_alerting_instance = MagicMock() proxy_logging.slack_alerting_instance.alert_types = [] proxy_logging.slack_alerting_instance.run_scheduled_deprecation_check = AsyncMock() proxy_logging._init_litellm_callbacks = MagicMock() proxy_logging.startup_event(llm_router=None, redis_usage_cache=None) assert proxy_logging.deprecation_check_started is False proxy_logging.update_values(alerting=["slack"]) assert proxy_logging.deprecation_check_started is True proxy_logging.slack_alerting_instance.run_scheduled_deprecation_check.assert_called_once_with( pod_lock_manager=proxy_logging.db_spend_update_writer.pod_lock_manager ) 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) @pytest.mark.asyncio async def test_startup_event_hands_the_daily_report_this_pods_lock_manager(proxy_logging): """regression: issue #14809 - the daily report's dedupe lock only works if startup_event passes the writer's pod_lock_manager down; dropping the argument silently restores the every-pod-reports behavior.""" proxy_logging.slack_alerting_instance = MagicMock() proxy_logging.slack_alerting_instance.alert_types = ["daily_reports"] proxy_logging.slack_alerting_instance._run_scheduled_daily_report = AsyncMock() proxy_logging._init_litellm_callbacks = MagicMock() proxy_logging.update_values = MagicMock() llm_router = MagicMock() proxy_logging.startup_event(llm_router=llm_router, redis_usage_cache=None) await asyncio.sleep(0) call = proxy_logging.slack_alerting_instance._run_scheduled_daily_report.call_args assert proxy_logging.slack_alerting_instance._run_scheduled_daily_report.call_count == 1 assert call.kwargs["pod_lock_manager"] is proxy_logging.db_spend_update_writer.pod_lock_manager assert call.kwargs["llm_router"] is llm_router @pytest.mark.asyncio async def test_startup_event_skips_the_daily_report_when_it_is_not_an_alert_type(proxy_logging): proxy_logging.slack_alerting_instance = MagicMock() proxy_logging.slack_alerting_instance.alert_types = [] proxy_logging.slack_alerting_instance._run_scheduled_daily_report = AsyncMock() proxy_logging._init_litellm_callbacks = MagicMock() proxy_logging.update_values = MagicMock() proxy_logging.startup_event(llm_router=None, redis_usage_cache=None) await asyncio.sleep(0) proxy_logging.slack_alerting_instance._run_scheduled_daily_report.assert_not_called() # --------------------------------------------------------------------------- # _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_iterations_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 _stub_hook_classes(): class _PrismaFreeHook: def __init__(self, internal_usage_cache): self.internal_usage_cache = internal_usage_cache class _PrismaRequiringHook: def __init__(self, internal_usage_cache, prisma_client): self.internal_usage_cache = internal_usage_cache self.prisma_client = prisma_client class _PrismaOnlyHook: def __init__(self, prisma_client): self.prisma_client = prisma_client return { "cache_control_check": _PrismaFreeHook, "needs_db_hook": _PrismaRequiringHook, "db_only_hook": _PrismaOnlyHook, } def test_add_proxy_hooks_skips_prisma_requiring_hook_when_no_db(proxy_logging, monkeypatch): hook_classes = _stub_hook_classes() registered: List[Any] = [] from litellm.proxy import utils as utils_mod monkeypatch.setattr(utils_mod, "PROXY_HOOKS", list(hook_classes.keys())) monkeypatch.setattr(utils_mod, "get_proxy_hook", hook_classes.__getitem__) 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) snapshot = { "mapping_keys": list(proxy_logging.proxy_hook_mapping.keys()), "registered_types": [type(r).__name__ for r in registered], "needs_db_hook_lookup": proxy_logging.get_proxy_hook("needs_db_hook"), "db_only_hook_lookup": proxy_logging.get_proxy_hook("db_only_hook"), } assert snapshot == { "mapping_keys": ["cache_control_check"], "registered_types": ["_PrismaFreeHook"], "needs_db_hook_lookup": None, "db_only_hook_lookup": None, } def test_add_proxy_hooks_registers_prisma_requiring_hook_with_db(proxy_logging, monkeypatch): hook_classes = _stub_hook_classes() registered: List[Any] = [] fake_prisma = MagicMock() from litellm.proxy import utils as utils_mod monkeypatch.setattr(utils_mod, "PROXY_HOOKS", list(hook_classes.keys())) monkeypatch.setattr(utils_mod, "get_proxy_hook", hook_classes.__getitem__) monkeypatch.setattr( litellm.logging_callback_manager, "add_litellm_callback", lambda cb: registered.append(cb), ) with patch("litellm.proxy.proxy_server.prisma_client", fake_prisma): proxy_logging._add_proxy_hooks(llm_router=None) snapshot = { "mapping_keys": list(proxy_logging.proxy_hook_mapping.keys()), "registered_count": len(registered), "needs_db_hook_got_prisma": proxy_logging.proxy_hook_mapping["needs_db_hook"].prisma_client is fake_prisma, "db_only_hook_got_prisma": proxy_logging.proxy_hook_mapping["db_only_hook"].prisma_client is fake_prisma, } assert snapshot == { "mapping_keys": ["cache_control_check", "needs_db_hook", "db_only_hook"], "registered_count": 3, "needs_db_hook_got_prisma": True, "db_only_hook_got_prisma": True, } 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_iterations = MagicMock() s_parallel = MagicMock() proxy_logging.proxy_hook_mapping = { "cache_control_check": s_cache, "max_iterations_limiter": s_iterations, "max_parallel_request_limiter": s_parallel, } snapshot = { "cache_control_check": proxy_logging.get_proxy_hook("cache_control_check") is s_cache, "max_iterations_limiter": proxy_logging.get_proxy_hook("max_iterations_limiter") is s_iterations, "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_iterations_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())