diff --git a/tests/test_litellm/enterprise/enterprise_callbacks/send_emails/test_base_email.py b/tests/test_litellm/enterprise/enterprise_callbacks/send_emails/test_base_email.py index fbf57afeeb4..744195dfb6f 100644 --- a/tests/test_litellm/enterprise/enterprise_callbacks/send_emails/test_base_email.py +++ b/tests/test_litellm/enterprise/enterprise_callbacks/send_emails/test_base_email.py @@ -20,6 +20,7 @@ from litellm_enterprise.types.enterprise_callbacks.send_emails import ( from litellm.integrations.email_templates.email_footer import EMAIL_FOOTER from litellm.proxy._types import CallInfo, Litellm_EntityType, WebhookEvent +from litellm.constants import EMAIL_BUDGET_ALERT_TTL @pytest.fixture(autouse=True) @@ -713,7 +714,7 @@ async def test_budget_alerts_soft_budget_crossed( cache_call_args = mock_cache.async_set_cache.call_args[1] assert cache_call_args["key"] == "email_budget_alerts:soft_budget_crossed:test_user" assert cache_call_args["value"] == "SENT" - assert cache_call_args["ttl"] == BaseEmailLogger.DEFAULT_BUDGET_ALERT_TTL + assert cache_call_args["ttl"] == EMAIL_BUDGET_ALERT_TTL @pytest.mark.asyncio @@ -839,4 +840,42 @@ async def test_get_email_params_soft_budget_crossed( # Should use default subject template for soft_budget_crossed assert result.subject == "LiteLLM: Soft Budget Crossed - Total Soft Budget: $100.0" assert result.recipient_email == "test@example.com" - assert result.base_url == "http://test.com" \ No newline at end of file + assert result.base_url == "http://test.com" + + +@pytest.mark.asyncio +async def test_budget_alerts_max_budget_alert_crossed( + base_email_logger, mock_send_email +): + """Test that budget_alerts sends email when max budget alert threshold is crossed""" + user_info = CallInfo( + user_id="test_user", + user_email="test@example.com", + spend=165.0, + max_budget=200.0, + event_group=Litellm_EntityType.USER, + ) + + mock_cache = mock.AsyncMock() + mock_cache.async_get_cache = mock.AsyncMock(return_value=None) + mock_cache.async_set_cache = mock.AsyncMock() + base_email_logger.internal_usage_cache = mock_cache + + with mock.patch.dict( + os.environ, + { + "PROXY_BASE_URL": "http://test.com", + }, + ): + await base_email_logger.budget_alerts(type="max_budget_alert", user_info=user_info) + + mock_send_email.assert_called_once() + call_args = mock_send_email.call_args[1] + assert call_args["to_email"] == ["test@example.com"] + assert "Max Budget Alert" in call_args["subject"] + + mock_cache.async_set_cache.assert_called_once() + cache_call_args = mock_cache.async_set_cache.call_args[1] + assert cache_call_args["key"] == "email_budget_alerts:max_budget_alert:test_user" + assert cache_call_args["value"] == "SENT" + assert cache_call_args["ttl"] == EMAIL_BUDGET_ALERT_TTL \ No newline at end of file diff --git a/tests/test_litellm/proxy/auth/test_auth_checks.py b/tests/test_litellm/proxy/auth/test_auth_checks.py index 52f1c433433..807559207e6 100644 --- a/tests/test_litellm/proxy/auth/test_auth_checks.py +++ b/tests/test_litellm/proxy/auth/test_auth_checks.py @@ -29,6 +29,7 @@ from litellm.proxy.auth.auth_checks import ( ExperimentalUIJWTToken, _can_object_call_vector_stores, _get_team_db_check, + _virtual_key_max_budget_alert_check, _virtual_key_soft_budget_check, get_user_object, vector_store_access_check, @@ -1132,3 +1133,147 @@ async def test_virtual_key_soft_budget_check_scenarios( assert ( alert_triggered == expect_alert ), f"Expected alert_triggered to be {expect_alert} for spend={spend}, soft_budget={soft_budget}" + + +@pytest.mark.asyncio +async def test_virtual_key_max_budget_alert_check_with_user_obj(): + """Test _virtual_key_max_budget_alert_check includes user_email when user_obj is provided""" + alert_triggered = False + captured_call_info = None + + class MockProxyLogging: + async def budget_alerts(self, type, user_info): + nonlocal alert_triggered, captured_call_info + alert_triggered = True + captured_call_info = user_info + assert type == "max_budget_alert" + assert isinstance(user_info, CallInfo) + + valid_token = UserAPIKeyAuth( + token="test-token", + spend=90.0, + max_budget=100.0, + user_id="test-user", + team_id="test-team", + team_alias="test-team-alias", + org_id="test-org", + key_alias="test-key", + soft_budget=50.0, + ) + + user_obj = LiteLLM_UserTable( + user_id="test-user", + user_email="test@example.com", + max_budget=None, + ) + + proxy_logging_obj = MockProxyLogging() + + await _virtual_key_max_budget_alert_check( + valid_token=valid_token, + proxy_logging_obj=proxy_logging_obj, + user_obj=user_obj, + ) + + await asyncio.sleep(0.1) + + assert alert_triggered is True + assert captured_call_info is not None + assert captured_call_info.user_email == "test@example.com" + assert captured_call_info.token == "test-token" + assert captured_call_info.spend == 90.0 + assert captured_call_info.max_budget == 100.0 + assert captured_call_info.soft_budget == 50.0 + assert captured_call_info.user_id == "test-user" + assert captured_call_info.team_id == "test-team" + assert captured_call_info.team_alias == "test-team-alias" + assert captured_call_info.organization_id == "test-org" + assert captured_call_info.key_alias == "test-key" + assert captured_call_info.event_group == Litellm_EntityType.KEY + + +@pytest.mark.asyncio +async def test_virtual_key_max_budget_alert_check_without_user_obj(): + """Test _virtual_key_max_budget_alert_check sets user_email to None when user_obj is not provided""" + alert_triggered = False + captured_call_info = None + + class MockProxyLogging: + async def budget_alerts(self, type, user_info): + nonlocal alert_triggered, captured_call_info + alert_triggered = True + captured_call_info = user_info + assert type == "max_budget_alert" + assert isinstance(user_info, CallInfo) + + valid_token = UserAPIKeyAuth( + token="test-token", + spend=90.0, + max_budget=100.0, + user_id="test-user", + team_id="test-team", + key_alias="test-key", + ) + + proxy_logging_obj = MockProxyLogging() + + await _virtual_key_max_budget_alert_check( + valid_token=valid_token, + proxy_logging_obj=proxy_logging_obj, + user_obj=None, + ) + + await asyncio.sleep(0.1) + + assert alert_triggered is True + assert captured_call_info is not None + assert captured_call_info.user_email is None + + +@pytest.mark.parametrize( + "spend, max_budget, expect_alert", + [ + (80.0, 100.0, True), # At 80% threshold (alert threshold) + (90.0, 100.0, True), # Above threshold, below max_budget + (79.0, 100.0, False), # Below threshold + (100.0, 100.0, False), # At max_budget (not below, so no alert) + (110.0, 100.0, False), # Above max_budget (already exceeded) + (100.0, None, False), # No max_budget set + (0.0, 100.0, False), # Spend is 0 + ], +) +@pytest.mark.asyncio +async def test_virtual_key_max_budget_alert_check_scenarios( + spend, max_budget, expect_alert +): + """Test _virtual_key_max_budget_alert_check with various spend and max_budget scenarios""" + alert_triggered = False + + class MockProxyLogging: + async def budget_alerts(self, type, user_info): + nonlocal alert_triggered + alert_triggered = True + assert type == "max_budget_alert" + assert isinstance(user_info, CallInfo) + + valid_token = UserAPIKeyAuth( + token="test-token", + spend=spend, + max_budget=max_budget, + user_id="test-user", + key_alias="test-key", + ) + + proxy_logging_obj = MockProxyLogging() + + await _virtual_key_max_budget_alert_check( + valid_token=valid_token, + proxy_logging_obj=proxy_logging_obj, + user_obj=None, + ) + + await asyncio.sleep(0.1) + + assert ( + alert_triggered == expect_alert + ), f"Expected alert_triggered to be {expect_alert} for spend={spend}, max_budget={max_budget}"