From 23a78f1f06f1bb9d89c8bbb4bd04bab49a997001 Mon Sep 17 00:00:00 2001
From: suleimanelkhoury <suleimanelkhoury@gmail.com>
Date: Tue, 14 Apr 2026 11:46:03 +0200
Subject: [PATCH] fix(tests): update budget enforcement tests to reflect new
 behavior in get_end_user_object

Since the budget enforcement is relocated from get_end_user_object inside auth_checks.py to  is_end_user_within_budget inside max_budget_limiter.py, both test cases test_get_end_user_object and test_budget_enforcement_blocks_over_budget_users are edited accordingly
---
 tests/proxy_unit_tests/test_auth_checks.py    | 40 +++++++------------
 .../test_default_end_user_budget_simple.py    | 22 +++++++---
 2 files changed, 31 insertions(+), 31 deletions(-)

diff --git a/tests/proxy_unit_tests/test_auth_checks.py b/tests/proxy_unit_tests/test_auth_checks.py
index c92ec61b9b2..c9389a38fd5 100644
--- a/tests/proxy_unit_tests/test_auth_checks.py
+++ b/tests/proxy_unit_tests/test_auth_checks.py
@@ -38,7 +38,7 @@ from litellm.proxy.utils import CallInfo
 async def test_get_end_user_object(customer_spend, customer_budget):
     """
     Scenario 1: normal
-    Scenario 2: user over budget
+    Scenario 2: user over budget (should still return object here, as the budget check happens in max_budget_limiter)
     """
     end_user_id = "my-test-customer"
     _budget = LiteLLM_BudgetTable(max_budget=customer_budget)
@@ -51,31 +51,19 @@ async def test_get_end_user_object(customer_spend, customer_budget):
     _cache = DualCache()
     _key = "end_user_id:{}".format(end_user_id)
     _cache.set_cache(key=_key, value=end_user_obj.model_dump())
-    try:
-        await get_end_user_object(
-            end_user_id=end_user_id,
-            prisma_client="RANDOM VALUE",  # type: ignore
-            user_api_key_cache=_cache,
-            route="/v1/chat/completions",
-        )
-        if customer_spend > customer_budget:
-            pytest.fail(
-                "Expected call to fail. Customer Spend={}, Customer Budget={}".format(
-                    customer_spend, customer_budget
-                )
-            )
-    except Exception as e:
-        if (
-            isinstance(e, litellm.BudgetExceededError)
-            and customer_spend > customer_budget
-        ):
-            pass
-        else:
-            pytest.fail(
-                "Expected call to work. Customer Spend={}, Customer Budget={}, Error={}".format(
-                    customer_spend, customer_budget, str(e)
-                )
-            )
+
+    result = await get_end_user_object(
+        end_user_id=end_user_id,
+        prisma_client="RANDOM VALUE",  # type: ignore
+        user_api_key_cache=_cache,
+        route="/v1/chat/completions",
+    )
+
+    assert result is not None
+    assert result.user_id == end_user_id
+    assert result.spend == customer_spend
+    assert result.litellm_budget_table is not None
+    assert result.litellm_budget_table.max_budget == customer_budget
 
 
 @pytest.mark.parametrize(
diff --git a/tests/proxy_unit_tests/test_default_end_user_budget_simple.py b/tests/proxy_unit_tests/test_default_end_user_budget_simple.py
index 92ca1f71703..21ae9b8da83 100644
--- a/tests/proxy_unit_tests/test_default_end_user_budget_simple.py
+++ b/tests/proxy_unit_tests/test_default_end_user_budget_simple.py
@@ -170,13 +170,25 @@ async def test_budget_enforcement_blocks_over_budget_users():
     mock_cache.async_get_cache = AsyncMock(return_value=None)
     mock_cache.async_set_cache = AsyncMock()
     
-    # Should raise BudgetExceededError
+    # 1. Fetch the user object (this no longer throws an exception)
+    result = await get_end_user_object(
+        end_user_id=end_user_id,
+        prisma_client=mock_prisma_client,
+        user_api_key_cache=mock_cache,
+        route="/chat/completions",
+    )
+    
+    from litellm.proxy.hooks.max_budget_limiter import _PROXY_MaxBudgetLimiter
+    budget_limiter = _PROXY_MaxBudgetLimiter()
+    assert default_budget.max_budget is not None
+
+    # 2. Check that the new limiter function correctly enforces the budget using the mocked values
     with pytest.raises(litellm.BudgetExceededError) as exc_info:
-        await get_end_user_object(
+        await budget_limiter.is_end_user_within_budget(
             end_user_id=end_user_id,
-            prisma_client=mock_prisma_client,
-            user_api_key_cache=mock_cache,
-            route="/chat/completions",
+            end_user_max_budget=default_budget.max_budget,
+            end_user_spend=mock_end_user_data["spend"],
+            route="/chat/completions"
         )
     
     assert "ExceededBudget" in str(exc_info.value)