diff --git a/litellm/tests/test_alerting.py b/litellm/tests/test_alerting.py index a74e25910c..5d6e068b7c 100644 --- a/litellm/tests/test_alerting.py +++ b/litellm/tests/test_alerting.py @@ -3,7 +3,7 @@ import sys import os -import io, asyncio, httpx +import io, asyncio from datetime import datetime, timedelta # import logging @@ -17,61 +17,6 @@ import asyncio from unittest.mock import patch, MagicMock from litellm.caching import DualCache from litellm.integrations.slack_alerting import SlackAlerting -from litellm.proxy._types import UserAPIKeyAuth -from litellm.proxy.proxy_server import HTTPException - - -@pytest.mark.parametrize("exception_type", ["llm-exception", "non-llm-exception"]) -@pytest.mark.asyncio -async def test_slack_alerting_llm_exceptions(exception_type, monkeypatch): - """ - Test if non-llm exception -> No request - Test if llm exception -> Request triggered - """ - _pl = ProxyLogging(user_api_key_cache=DualCache()) - _pl.update_values( - alerting=["slack"], - alerting_threshold=100, - redis_cache=None, - alert_types=["llm_exceptions"], - ) - - async def mock_alerting_handler(message, level, alert_type): - global exception_type - - if exception_type == "llm-exception": - pass - elif exception_type == "non-llm-exception": - pytest.fail("Function should not have been called") - - monkeypatch.setattr(_pl, "alerting_handler", mock_alerting_handler) - - if exception_type == "llm-exception": - await _pl.post_call_failure_hook( - original_exception=litellm.APIError( - status_code=500, - message="This is a test exception", - llm_provider="openai", - model="gpt-3.5-turbo", - request=httpx.Request( - method="completion", url="https://github.com/BerriAI/litellm" - ), - ), - user_api_key_dict=UserAPIKeyAuth(), - ) - - await asyncio.sleep(2) - - elif exception_type == "non-llm-exception": - await _pl.post_call_failure_hook( - original_exception=HTTPException( - status_code=400, - detail={"error": "this is a test exception"}, - ), - user_api_key_dict=UserAPIKeyAuth(), - ) - - await asyncio.sleep(2) @pytest.mark.asyncio @@ -149,3 +94,80 @@ def test_init(): assert slack_no_alerting.alerting == [] print("passed testing slack alerting init") + + +from unittest.mock import patch, AsyncMock +from datetime import datetime, timedelta + + +@pytest.fixture +def slack_alerting(): + return SlackAlerting(alerting_threshold=1) + + +# Test for hanging LLM responses +@pytest.mark.asyncio +async def test_response_taking_too_long_hanging(slack_alerting): + request_data = { + "model": "test_model", + "messages": "test_messages", + "litellm_status": "running", + } + with patch.object(slack_alerting, "send_alert", new=AsyncMock()) as mock_send_alert: + await slack_alerting.response_taking_too_long( + type="hanging_request", request_data=request_data + ) + mock_send_alert.assert_awaited_once() + + +# Test for slow LLM responses +@pytest.mark.asyncio +async def test_response_taking_too_long_callback(slack_alerting): + start_time = datetime.now() + end_time = start_time + timedelta(seconds=301) + kwargs = {"model": "test_model", "messages": "test_messages", "litellm_params": {}} + with patch.object(slack_alerting, "send_alert", new=AsyncMock()) as mock_send_alert: + await slack_alerting.response_taking_too_long_callback( + kwargs, None, start_time, end_time + ) + mock_send_alert.assert_awaited_once() + + +# Test for budget crossed +@pytest.mark.asyncio +async def test_budget_alerts_crossed(slack_alerting): + user_max_budget = 100 + user_current_spend = 101 + with patch.object(slack_alerting, "send_alert", new=AsyncMock()) as mock_send_alert: + await slack_alerting.budget_alerts( + "user_budget", user_max_budget, user_current_spend + ) + mock_send_alert.assert_awaited_once() + + +# Test for budget crossed again (should not fire alert 2nd time) +@pytest.mark.asyncio +async def test_budget_alerts_crossed_again(slack_alerting): + user_max_budget = 100 + user_current_spend = 101 + with patch.object(slack_alerting, "send_alert", new=AsyncMock()) as mock_send_alert: + await slack_alerting.budget_alerts( + "user_budget", user_max_budget, user_current_spend + ) + mock_send_alert.assert_awaited_once() + mock_send_alert.reset_mock() + await slack_alerting.budget_alerts( + "user_budget", user_max_budget, user_current_spend + ) + mock_send_alert.assert_not_awaited() + + +# Test for send_alert - should be called once +@pytest.mark.asyncio +async def test_send_alert(slack_alerting): + with patch.object( + slack_alerting.async_http_handler, "post", new=AsyncMock() + ) as mock_post: + mock_post.return_value.status_code = 200 + await slack_alerting.send_alert("Test message", "Low", "budget_alerts") + mock_post.assert_awaited_once()