From 20175969130bdeccd5ce643ed3d9e58b30d4fe31 Mon Sep 17 00:00:00 2001
From: Ishaan Jaff <ishaanjaffer0324@gmail.com>
Date: Fri, 24 Jan 2025 15:32:41 -0800
Subject: [PATCH] Revert "test_team_and_key_budget_enforcement"

This reverts commit 9d44f51847998286020d6820daf2472d522fcd19.
---
 litellm/proxy/auth/auth_checks.py      |   2 +-
 tests/otel_tests/test_e2e_budgeting.py | 108 -------------------------
 2 files changed, 1 insertion(+), 109 deletions(-)

diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py
index baca826804..317c7fddc4 100644
--- a/litellm/proxy/auth/auth_checks.py
+++ b/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"Budget has been exceeded! Team={team_object.team_id} 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}",
         )
 
 
diff --git a/tests/otel_tests/test_e2e_budgeting.py b/tests/otel_tests/test_e2e_budgeting.py
index 5986f3f013..227d049f82 100644
--- a/tests/otel_tests/test_e2e_budgeting.py
+++ b/tests/otel_tests/test_e2e_budgeting.py
@@ -3,7 +3,6 @@ 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):
@@ -279,110 +278,3 @@ 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"