"""Pin alerting helpers on ``ProxyLogging``. Covers ``failed_tracking_alert``, ``budget_alerts``, ``alerting_handler``, ``failure_handler``. """ from __future__ import annotations from datetime import datetime from typing import Any, Dict, Final from unittest.mock import AsyncMock, MagicMock import pytest from fastapi import HTTPException from prisma.errors import PrismaError from litellm._service_logger import ServiceTypes from litellm.litellm_core_utils import litellm_logging from litellm.proxy._types import AlertType, CallInfo # --------------------------------------------------------------------------- # failed_tracking_alert # --------------------------------------------------------------------------- @pytest.mark.asyncio async def test_failed_tracking_alert_no_op_when_alerting_none(proxy_logging): proxy_logging.alerting = None proxy_logging.slack_alerting_instance = MagicMock(failed_tracking_alert=AsyncMock()) await proxy_logging.failed_tracking_alert(error_message="x", failing_model="m") proxy_logging.slack_alerting_instance.failed_tracking_alert.assert_not_called() @pytest.mark.asyncio async def test_failed_tracking_alert_forwards_to_slack(proxy_logging): proxy_logging.alerting = ["slack"] captured: Dict[str, Any] = {} async def fake_alert(**kwargs): captured.update(kwargs) proxy_logging.slack_alerting_instance = MagicMock(failed_tracking_alert=fake_alert) await proxy_logging.failed_tracking_alert(error_message="db down", failing_model="gpt-4") snapshot = { "error_message": captured["error_message"], "failing_model": captured["failing_model"], "captured_keys": sorted(captured.keys()), } assert snapshot == { "error_message": "db down", "failing_model": "gpt-4", "captured_keys": ["error_message", "failing_model"], } @pytest.mark.asyncio async def test_failed_tracking_alert_slack_error_raises(proxy_logging): proxy_logging.alerting = ["slack"] proxy_logging.slack_alerting_instance = MagicMock( failed_tracking_alert=AsyncMock(side_effect=RuntimeError("slack down")) ) with pytest.raises(RuntimeError): await proxy_logging.failed_tracking_alert(error_message="x", failing_model="m") # --------------------------------------------------------------------------- # budget_alerts # --------------------------------------------------------------------------- def _user_info(alert_emails=None): return CallInfo( spend=0.0, max_budget=1.0, token="tok", user_id="u1", team_id="t1", team_alias=None, user_email=None, key_alias=None, projected_exceeded_date=None, projected_spend=None, event_group="user", event="threshold_crossed", alert_emails=alert_emails, ) @pytest.mark.asyncio async def test_budget_alerts_no_op_when_alerting_off_and_no_emails(proxy_logging): proxy_logging.alerting = None proxy_logging.slack_alerting_instance = MagicMock(budget_alerts=AsyncMock()) proxy_logging.email_logging_instance = MagicMock(budget_alerts=AsyncMock()) await proxy_logging.budget_alerts(type="user_budget", user_info=_user_info()) proxy_logging.slack_alerting_instance.budget_alerts.assert_not_called() proxy_logging.email_logging_instance.budget_alerts.assert_not_called() @pytest.mark.asyncio async def test_budget_alerts_slack_when_slack_alerting(proxy_logging): proxy_logging.alerting = ["slack"] captured: Dict[str, Any] = {} async def fake_alert(**kwargs): captured.update(kwargs) proxy_logging.slack_alerting_instance = MagicMock(budget_alerts=fake_alert) proxy_logging.email_logging_instance = None user_info = _user_info() await proxy_logging.budget_alerts(type="user_budget", user_info=user_info) snapshot = { "type": captured["type"], "user_info_is_callinfo": isinstance(captured["user_info"], CallInfo), "user_id": captured["user_info"].user_id, } assert snapshot == {"type": "user_budget", "user_info_is_callinfo": True, "user_id": "u1"} @pytest.mark.asyncio async def test_budget_alerts_webhook_only_forwards_to_slack_alerting_instance(proxy_logging): proxy_logging.alerting = ["webhook"] captured: Dict[str, Any] = {} async def fake_alert(**kwargs): captured.update(kwargs) proxy_logging.slack_alerting_instance = MagicMock(budget_alerts=fake_alert) proxy_logging.email_logging_instance = None await proxy_logging.budget_alerts(type="user_budget", user_info=_user_info()) snapshot = { "type": captured["type"], "user_info_is_callinfo": isinstance(captured["user_info"], CallInfo), "user_id": captured["user_info"].user_id, } assert snapshot == {"type": "user_budget", "user_info_is_callinfo": True, "user_id": "u1"} @pytest.mark.asyncio async def test_budget_alerts_email_only_skips_slack_alerting_instance(proxy_logging): proxy_logging.alerting = ["email"] proxy_logging.slack_alerting_instance = MagicMock(budget_alerts=AsyncMock()) proxy_logging.email_logging_instance = MagicMock(budget_alerts=AsyncMock()) await proxy_logging.budget_alerts(type="user_budget", user_info=_user_info()) proxy_logging.slack_alerting_instance.budget_alerts.assert_not_called() proxy_logging.email_logging_instance.budget_alerts.assert_called_once() @pytest.mark.asyncio async def test_budget_alerts_soft_budget_with_alert_emails_bypasses_global(proxy_logging): proxy_logging.alerting = None proxy_logging.slack_alerting_instance = MagicMock(budget_alerts=AsyncMock()) proxy_logging.email_logging_instance = MagicMock(budget_alerts=AsyncMock()) info = _user_info(alert_emails=["a@b.c"]) await proxy_logging.budget_alerts(type="soft_budget", user_info=info) proxy_logging.email_logging_instance.budget_alerts.assert_called_once() proxy_logging.slack_alerting_instance.budget_alerts.assert_not_called() @pytest.mark.asyncio async def test_budget_alerts_slack_failure_raises(proxy_logging): proxy_logging.alerting = ["slack"] proxy_logging.slack_alerting_instance = MagicMock( budget_alerts=AsyncMock(side_effect=ConnectionError("slack")) ) proxy_logging.email_logging_instance = None with pytest.raises(ConnectionError): await proxy_logging.budget_alerts(type="user_budget", user_info=_user_info()) # --------------------------------------------------------------------------- # alerting_handler # --------------------------------------------------------------------------- @pytest.mark.asyncio async def test_alerting_handler_no_op_when_alerting_is_none(proxy_logging): proxy_logging.alerting = None proxy_logging.slack_alerting_instance = MagicMock(send_alert=AsyncMock()) await proxy_logging.alerting_handler(message="x", level="High", alert_type=AlertType.db_exceptions) proxy_logging.slack_alerting_instance.send_alert.assert_not_called() @pytest.mark.asyncio async def test_alerting_handler_sends_to_slack(proxy_logging): proxy_logging.alerting = ["slack"] captured: Dict[str, Any] = {} async def fake_send(**kwargs): captured.update(kwargs) proxy_logging.slack_alerting_instance = MagicMock(send_alert=fake_send) await proxy_logging.alerting_handler( message="hi", level="High", alert_type=AlertType.db_exceptions, request_data={"metadata": {}} ) snapshot = { "message": captured["message"], "level": captured["level"], "alert_type": captured["alert_type"], "user_info": captured["user_info"], } assert snapshot == { "message": "hi", "level": "High", "alert_type": AlertType.db_exceptions, "user_info": None, } @pytest.mark.asyncio async def test_alerting_handler_sentry_without_sdk_error_raises(proxy_logging, monkeypatch): proxy_logging.alerting = ["sentry"] monkeypatch.delenv("SENTRY_DSN", raising=False) monkeypatch.setattr(litellm_logging, "sentry_sdk_instance", None) litellm_logging.start_sentry.cache_clear() with pytest.raises(Exception, match="SENTRY_DSN"): await proxy_logging.alerting_handler(message="x", level="Low", alert_type=AlertType.db_exceptions) @pytest.mark.asyncio async def test_alerting_handler_sentry_sends_through_the_initialized_sdk(proxy_logging, monkeypatch): proxy_logging.alerting = ["sentry"] sentry = MagicMock() monkeypatch.setattr(litellm_logging, "sentry_sdk_instance", sentry) await proxy_logging.alerting_handler(message="db down", level="High", alert_type=AlertType.db_exceptions) sentry.capture_message.assert_called_once() assert "db down" in sentry.capture_message.call_args.args[0] @pytest.mark.asyncio async def test_alerting_handler_starts_sentry_when_the_alert_was_not_configured_at_boot(proxy_logging, monkeypatch): import sentry_sdk monkeypatch.setattr(proxy_logging, "alerting", ["sentry"]) init: Final = MagicMock() capture_message: Final = MagicMock() monkeypatch.setattr(sentry_sdk, "init", init) monkeypatch.setattr(sentry_sdk, "capture_message", capture_message) monkeypatch.setenv("SENTRY_DSN", "https://key@sentry.example/1") monkeypatch.setattr(litellm_logging, "sentry_sdk_instance", None) litellm_logging.start_sentry.cache_clear() monkeypatch.setattr(litellm_logging, "capture_exception", None) monkeypatch.setattr(litellm_logging, "add_breadcrumb", None) await proxy_logging.alerting_handler(message="db down", level="High", alert_type=AlertType.db_exceptions) assert init.call_args.kwargs["dsn"] == "https://key@sentry.example/1" capture_message.assert_called_once() assert "db down" in capture_message.call_args.args[0] # --------------------------------------------------------------------------- # failure_handler # --------------------------------------------------------------------------- @pytest.mark.asyncio async def test_failure_handler_skips_when_db_exceptions_not_in_alert_types(proxy_logging): proxy_logging.alert_types = ["llm_too_slow"] # type: ignore[list-item] proxy_logging.alerting_handler = AsyncMock() proxy_logging.service_logging_obj = MagicMock(async_service_failure_hook=AsyncMock()) await proxy_logging.failure_handler(original_exception=Exception("x"), duration=1.0, call_type="db_read") proxy_logging.alerting_handler.assert_not_called() proxy_logging.service_logging_obj.async_service_failure_hook.assert_not_called() @pytest.mark.asyncio async def test_failure_handler_logs_db_error_and_calls_service_logging(proxy_logging, monkeypatch): proxy_logging.alert_types = [AlertType.db_exceptions] proxy_logging.alerting_handler = AsyncMock() proxy_logging.service_logging_obj = MagicMock(async_service_failure_hook=AsyncMock()) monkeypatch.setattr(litellm_logging, "capture_exception", None) await proxy_logging.failure_handler( original_exception=HTTPException(status_code=500, detail="boom"), duration=1.5, call_type="db_write", ) call_kwargs = proxy_logging.service_logging_obj.async_service_failure_hook.call_args.kwargs snapshot = { "service": call_kwargs["service"].value if hasattr(call_kwargs["service"], "value") else call_kwargs["service"], "duration": call_kwargs["duration"], "call_type": call_kwargs["call_type"], } assert snapshot == { "service": "postgres", "duration": 1.5, "call_type": "db_write", } @pytest.mark.asyncio @pytest.mark.parametrize("call_type", ["get_data", "insert_data", "update_data", "delete_data"]) async def test_failure_handler_alerts_but_leaves_prisma_error_event_to_log_db_metrics( proxy_logging, monkeypatch, call_type ): proxy_logging.alert_types = [AlertType.db_exceptions] proxy_logging.alerting_handler = AsyncMock() proxy_logging.service_logging_obj = MagicMock(async_service_failure_hook=AsyncMock()) monkeypatch.setattr(litellm_logging, "capture_exception", None) await proxy_logging.failure_handler( original_exception=PrismaError("connection reset"), duration=1.0, call_type=call_type ) snapshot = { "alerting_handler_scheduled": proxy_logging.alerting_handler.called, "service_failure_called": proxy_logging.service_logging_obj.async_service_failure_hook.called, } assert snapshot == {"alerting_handler_scheduled": True, "service_failure_called": False} @pytest.mark.asyncio async def test_failure_handler_still_emits_db_event_for_wrapped_insert_error(proxy_logging, monkeypatch): proxy_logging.alert_types = [AlertType.db_exceptions] proxy_logging.alerting_handler = AsyncMock() proxy_logging.service_logging_obj = MagicMock(async_service_failure_hook=AsyncMock()) monkeypatch.setattr(litellm_logging, "capture_exception", None) await proxy_logging.failure_handler( original_exception=HTTPException(status_code=400, detail={"error": "Foreign Key Constraint failed"}), duration=1.0, call_type="insert_data", ) call_kwargs = proxy_logging.service_logging_obj.async_service_failure_hook.call_args.kwargs assert (call_kwargs["service"], call_kwargs["call_type"]) == (ServiceTypes.DB, "insert_data") @pytest.mark.asyncio async def test_failure_handler_with_capture_exception_invoked(proxy_logging, monkeypatch): proxy_logging.alert_types = [AlertType.db_exceptions] proxy_logging.alerting_handler = AsyncMock() proxy_logging.service_logging_obj = MagicMock(async_service_failure_hook=AsyncMock()) captured: Dict[str, Any] = {} def fake_capture(error): captured["error"] = error monkeypatch.setattr(litellm_logging, "capture_exception", fake_capture) err = RuntimeError("real") await proxy_logging.failure_handler(original_exception=err, duration=1.0, call_type="db_read") snapshot = { "captured_is_input": captured["error"] is err, "service_failure_called": proxy_logging.service_logging_obj.async_service_failure_hook.called, "alerting_handler_scheduled": proxy_logging.alerting_handler.called, } assert snapshot == { "captured_is_input": True, "service_failure_called": True, "alerting_handler_scheduled": True, } @pytest.mark.asyncio async def test_failure_handler_propagates_service_logging_error_raises(proxy_logging, monkeypatch): proxy_logging.alert_types = [AlertType.db_exceptions] proxy_logging.alerting_handler = AsyncMock() proxy_logging.service_logging_obj = MagicMock( async_service_failure_hook=AsyncMock(side_effect=RuntimeError("svc")) ) monkeypatch.setattr(litellm_logging, "capture_exception", None) with pytest.raises(RuntimeError): await proxy_logging.failure_handler( original_exception=Exception("x"), duration=0.0, call_type="db_read" )