(Testing) e2e testing for team budget enforcement checks (#7988)

* test_team_and_key_budget_enforcement

* test_team_budget_update

* test_gemini_pro_json_schema_httpx_content_policy_error

litellm/proxy/auth/auth_checks.py

@@ -1037,7 +1037,7 @@ async def _team_max_budget_check(
         raise litellm.BudgetExceededError(
             current_cost=team_object.spend,
             max_budget=team_object.max_budget,
-            message=f"Team={team_object.team_id} over budget. Spend={team_object.spend}, Budget={team_object.max_budget}",
+            message=f"Budget has been exceeded! Team={team_object.team_id} Current cost: {team_object.spend}, Max budget: {team_object.max_budget}",
         )
 
 

tests/local_testing/test_amazing_vertex_completion.py

@@ -17,7 +17,7 @@ import asyncio
 import json
 import os
 import tempfile
-from unittest.mock import AsyncMock, MagicMock, patch
+from unittest.mock import AsyncMock, MagicMock, patch, ANY
 from respx import MockRouter
 import httpx
 
@@ -1222,6 +1222,7 @@ Using this JSON schema:
             messages=messages,
             response_format={"type": "json_object"},
             client=client,
+            logging_obj=ANY,
         )
 
         assert response.choices[0].finish_reason == "content_filter"

tests/otel_tests/test_e2e_budgeting.py

@@ -3,6 +3,7 @@ import asyncio
 import aiohttp
 import json
 from httpx import AsyncClient
+from typing import Any, Optional
 
 
 async def make_calls_until_budget_exceeded(session, key: str, call_function, **kwargs):
@@ -278,3 +279,174 @@ async def test_team_limit_modifications(field):
     print("response: ", json.dumps(response.json(), indent=4))
     assert response.status_code == 200
     assert response.json()["data"][field] is None
+
+
+async def generate_team_key(
+    session,
+    team_id: str,
+    max_budget: Optional[float] = None,
+):
+    """Helper function to generate a key for a specific team"""
+    url = "http://0.0.0.0:4000/key/generate"
+    headers = {"Authorization": "Bearer sk-1234", "Content-Type": "application/json"}
+    data: dict[str, Any] = {"team_id": team_id}
+    if max_budget is not None:
+        data["max_budget"] = max_budget
+    async with session.post(url, headers=headers, json=data) as response:
+        return await response.json()
+
+
+async def create_team(
+    session,
+    max_budget=None,
+):
+    """Helper function to create a new team"""
+    url = "http://0.0.0.0:4000/team/new"
+    headers = {"Authorization": "Bearer sk-1234", "Content-Type": "application/json"}
+    data = {
+        "max_budget": max_budget,
+    }
+    async with session.post(url, headers=headers, json=data) as response:
+        return await response.json()
+
+
+@pytest.mark.asyncio
+async def test_team_budget_enforcement():
+    """
+    Test budget enforcement for team-wide budgets:
+    1. Create team with low budget
+    2. Create key for that team
+    3. Make calls until team budget exceeded
+    4. Verify budget exceeded error
+    """
+    async with aiohttp.ClientSession() as session:
+        # Create team with low budget
+        team_response = await create_team(session=session, max_budget=0.0000000005)
+        team_id = team_response["team_id"]
+
+        # Create key for team (no specific budget)
+        key_gen = await generate_team_key(session=session, team_id=team_id)
+        key = key_gen["key"]
+
+        # Make calls until budget exceeded
+        calls_made = await make_calls_until_budget_exceeded(
+            session=session,
+            key=key,
+            call_function=chat_completion,
+            model="fake-openai-endpoint",
+        )
+
+        assert (
+            calls_made > 0
+        ), "Should make at least one successful call before team budget exceeded"
+
+
+@pytest.mark.asyncio
+async def test_team_and_key_budget_enforcement():
+    """
+    Test budget enforcement when both team and key have budgets:
+    1. Create team with low budget
+    2. Create key with higher budget
+    3. Verify team budget is enforced first
+    """
+    async with aiohttp.ClientSession() as session:
+        # Create team with very low budget
+        team_response = await create_team(session=session, max_budget=0.0000000005)
+        team_id = team_response["team_id"]
+
+        # Create key with higher budget
+        key_gen = await generate_team_key(
+            session=session,
+            team_id=team_id,
+            max_budget=0.001,  # Higher than team budget
+        )
+        key = key_gen["key"]
+
+        # Make calls until budget exceeded
+        calls_made = await make_calls_until_budget_exceeded(
+            session=session,
+            key=key,
+            call_function=chat_completion,
+            model="fake-openai-endpoint",
+        )
+
+        assert (
+            calls_made > 0
+        ), "Should make at least one successful call before team budget exceeded"
+
+        # Verify it was the team budget that was exceeded
+        try:
+            await chat_completion(
+                session=session, key=key, model="fake-openai-endpoint"
+            )
+        except Exception as e:
+            error_dict = e.body
+            assert (
+                "Budget has been exceeded! Team=" in error_dict["message"]
+            ), "Error should mention team budget being exceeded"
+
+            assert team_id in error_dict["message"], "Error should mention team id"
+
+
+async def update_team_budget(session, team_id: str, max_budget: float):
+    """Helper function to update a team's max budget"""
+    url = "http://0.0.0.0:4000/team/update"
+    headers = {"Authorization": "Bearer sk-1234", "Content-Type": "application/json"}
+    data = {
+        "team_id": team_id,
+        "max_budget": max_budget,
+    }
+    async with session.post(url, headers=headers, json=data) as response:
+        return await response.json()
+
+
+@pytest.mark.asyncio
+async def test_team_budget_update():
+    """
+    Test that requests continue working after updating a team's budget:
+    1. Create team with low budget
+    2. Create key for that team
+    3. Make calls until team budget exceeded
+    4. Update team with higher budget
+    5. Verify calls work again
+    """
+    async with aiohttp.ClientSession() as session:
+        # Create team with very low budget
+        team_response = await create_team(session=session, max_budget=0.0000000005)
+        team_id = team_response["team_id"]
+
+        # Create key for team (no specific budget)
+        key_gen = await generate_team_key(session=session, team_id=team_id)
+        key = key_gen["key"]
+
+        # Make calls until budget exceeded
+        calls_made = await make_calls_until_budget_exceeded(
+            session=session,
+            key=key,
+            call_function=chat_completion,
+            model="fake-openai-endpoint",
+        )
+
+        assert (
+            calls_made > 0
+        ), "Should make at least one successful call before team budget exceeded"
+
+        # Update team with higher budget
+        await update_team_budget(session, team_id, max_budget=0.001)
+
+        # Verify calls work again
+        for _ in range(3):
+            try:
+                response = await chat_completion(
+                    session=session, key=key, model="fake-openai-endpoint"
+                )
+                print("response: ", response)
+                assert (
+                    response is not None
+                ), "Should get valid response after budget update"
+            except Exception as e:
+                pytest.fail(
+                    f"Request should succeed after team budget update but got error: {e}"
+                )
+
+        # Verify it was the team budget that was exceeded