Tests

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"
+        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

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}"