From e62ada9d63807442955f4c873696eddf8f07b369 Mon Sep 17 00:00:00 2001
From: Sricharan Anugandula <anugandulasricharan99@gmail.com>
Date: Wed, 9 Sep 2026 01:13:41 +0530
Subject: [PATCH] added test cases

---
 .../auth/test_custom_auth_end_user_budget.py  | 84 +++++++++++++++++++
 1 file changed, 84 insertions(+)

diff --git a/tests/test_litellm/proxy/auth/test_custom_auth_end_user_budget.py b/tests/test_litellm/proxy/auth/test_custom_auth_end_user_budget.py
index cf1f665ad21..baa7e63ffe2 100644
--- a/tests/test_litellm/proxy/auth/test_custom_auth_end_user_budget.py
+++ b/tests/test_litellm/proxy/auth/test_custom_auth_end_user_budget.py
@@ -277,3 +277,87 @@ def test_update_valid_token_db_values_override_custom_auth_when_set():
     # DB values should win
     assert result.end_user_tpm_limit == 500
     assert result.end_user_model_max_budget == db_budget
+
+
+def test_update_valid_token_copies_end_user_max_budget():
+    # Regression: end_user_max_budget was the only limit field not copied by
+    # update_valid_token_with_end_user_params, so a budget from
+    # max_end_user_budget_id was silently dropped and budget reservation
+    # returned None, letting all concurrent first requests bypass the limit.
+    valid_token = UserAPIKeyAuth(token="test_token", end_user_id="customer-1")
+
+    end_user_params = {
+        "end_user_id": "customer-1",
+        "end_user_max_budget": 0.000000001,
+    }
+
+    result = update_valid_token_with_end_user_params(valid_token, end_user_params)
+
+    assert result.end_user_max_budget == 0.000000001
+
+
+def test_update_valid_token_preserves_custom_auth_max_budget_when_db_has_none():
+    # If custom auth sets end_user_max_budget and the DB budget table has no
+    # max_budget, the custom-auth-supplied value must survive the copy.
+    valid_token = UserAPIKeyAuth(
+        token="test_token",
+        end_user_id="customer-1",
+        end_user_max_budget=50.0,
+    )
+
+    end_user_params = {
+        "end_user_id": "customer-1",
+        # no end_user_max_budget key: DB found no value
+    }
+
+    result = update_valid_token_with_end_user_params(valid_token, end_user_params)
+
+    assert result.end_user_max_budget == 50.0
+
+
+@pytest.mark.asyncio
+async def test_end_user_budget_counter_created_from_token_max_budget():
+    # Regression: _get_end_user_budget_counter reads valid_token.end_user_max_budget
+    # first. When that field was always None (due to the copy bug), no counter was
+    # returned and budget reservation skipped enforcement entirely for new end users.
+    from litellm.proxy.spend_tracking.budget_reservation import (
+        _get_end_user_budget_counter,
+    )
+
+    token = UserAPIKeyAuth(
+        token="test_token",
+        end_user_id="customer-1",
+        end_user_max_budget=0.000000001,
+    )
+
+    counter = await _get_end_user_budget_counter(
+        valid_token=token,
+        end_user_id="customer-1",
+        end_user_object=None,
+    )
+
+    assert counter is not None
+    assert counter.max_budget == 0.000000001
+    assert counter.counter_key == "spend:end_user:customer-1"
+
+
+@pytest.mark.asyncio
+async def test_end_user_budget_counter_none_when_max_budget_missing():
+    # Without a budget on the token or the end_user_object, no counter should
+    # be returned: the user is unrestricted and reservation is skipped correctly.
+    from litellm.proxy.spend_tracking.budget_reservation import (
+        _get_end_user_budget_counter,
+    )
+
+    token = UserAPIKeyAuth(
+        token="test_token",
+        end_user_id="customer-1",
+    )
+
+    counter = await _get_end_user_budget_counter(
+        valid_token=token,
+        end_user_id="customer-1",
+        end_user_object=None,
+    )
+
+    assert counter is None