From a8fba14e107ccaad5e972895606f2efeb53efab0 Mon Sep 17 00:00:00 2001
From: Yuneng Jiang <yuneng@berri.ai>
Date: Tue, 15 Sep 2026 19:48:25 -0700
Subject: [PATCH] test(proxy): assert budget resets decrement the cleared spend

The reset job moved from zeroing spend to an atomic decrement of the amount
it cleared, so every batched write now carries {"decrement": <cleared>}
instead of 0. Four tests still pinned 0 and had been failing since, which
also meant they no longer checked the amount at all. Assert the decrement
equals each row's own pre-reset spend, so a wrong amount fails the test.
---
 .../test_proxy_budget_reset.py                | 24 +++++++++++++++----
 1 file changed, 20 insertions(+), 4 deletions(-)

diff --git a/tests/litellm_utils_tests/test_proxy_budget_reset.py b/tests/litellm_utils_tests/test_proxy_budget_reset.py
index 4103536950d..fe3c38a771f 100644
--- a/tests/litellm_utils_tests/test_proxy_budget_reset.py
+++ b/tests/litellm_utils_tests/test_proxy_budget_reset.py
@@ -163,6 +163,9 @@ async def test_reset_budget_keys_partial_failure():
     key1, key2, key3, key4, key5, key6 = (
         _attrify(k) for k in [key1, key2, key3, key4, key5, key6]
     )
+    pre_reset_spend = {
+        k["token"]: k["spend"] for k in [key2, key3, key4, key5, key6]
+    }
     prisma_client.get_data = AsyncMock(
         return_value=[key1, key2, key3, key4, key5, key6]
     )
@@ -201,7 +204,7 @@ async def test_reset_budget_keys_partial_failure():
     # And every write must carry only {spend, budget_reset_at} — never the full row.
     for c in key_writes:
         assert set(c["data"].keys()) == {"spend", "budget_reset_at"}
-        assert c["data"]["spend"] == 0
+        assert c["data"]["spend"] == {"decrement": pre_reset_spend[c["where"]["token"]]}
 
     # Verify that the failure logging hook was scheduled (due to the failure for key1)
     failure_hook_calls = (
@@ -252,6 +255,9 @@ async def test_reset_budget_users_partial_failure():
     user1, user2, user3, user4, user5, user6 = (
         _attrify(u) for u in [user1, user2, user3, user4, user5, user6]
     )
+    pre_reset_spend = {
+        u["user_id"]: u["spend"] for u in [user2, user3, user4, user5, user6]
+    }
     prisma_client.get_data = AsyncMock(
         return_value=[user1, user2, user3, user4, user5, user6]
     )
@@ -280,7 +286,9 @@ async def test_reset_budget_users_partial_failure():
     assert written_ids == ["user2", "user3", "user4", "user5", "user6"]
     for c in user_writes:
         assert set(c["data"].keys()) == {"spend", "budget_reset_at"}
-        assert c["data"]["spend"] == 0
+        assert c["data"]["spend"] == {
+            "decrement": pre_reset_spend[c["where"]["user_id"]]
+        }
 
     failure_hook_calls = (
         proxy_logging_obj.service_logging_obj.async_service_failure_hook.call_args_list
@@ -441,6 +449,7 @@ async def test_reset_budget_teams_partial_failure():
     for t in [team1, team2]:
         t.setdefault("team_id", t["id"])
     team1, team2 = _attrify(team1), _attrify(team2)
+    pre_reset_spend = team2["spend"]
     prisma_client.get_data = AsyncMock(return_value=[team1, team2])
 
     async def fake_reset_team(team, current_time, reset_settings=None):
@@ -465,7 +474,7 @@ async def test_reset_budget_teams_partial_failure():
     assert len(team_writes) == 1
     assert team_writes[0]["where"] == {"team_id": "team2"}
     assert set(team_writes[0]["data"].keys()) == {"spend", "budget_reset_at"}
-    assert team_writes[0]["data"]["spend"] == 0
+    assert team_writes[0]["data"]["spend"] == {"decrement": pre_reset_spend}
 
     failure_hook_calls = (
         proxy_logging_obj.service_logging_obj.async_service_failure_hook.call_args_list
@@ -542,6 +551,11 @@ async def test_reset_budget_continues_other_categories_on_failure():
     user1, user2 = _attrify(user1), _attrify(user2)
     team1, team2 = _attrify(team1), _attrify(team2)
     enduser1 = _attrify(enduser1)
+    pre_reset_spend = {
+        **{k["token"]: k["spend"] for k in [key1, key2]},
+        **{u["user_id"]: u["spend"] for u in [user2]},
+        **{t["team_id"]: t["spend"] for t in [team1, team2]},
+    }
     _wire_cascade_reads_for_test(prisma_client)
 
     proxy_logging_obj = MagicMock()
@@ -618,7 +632,9 @@ async def test_reset_budget_continues_other_categories_on_failure():
     # Every batched write must carry only the two reset fields, never the full row.
     for c in key_writes + user_writes + team_writes:
         assert set(c["data"].keys()) == {"spend", "budget_reset_at"}
-        assert c["data"]["spend"] == 0
+        assert c["data"]["spend"] == {
+            "decrement": pre_reset_spend[next(iter(c["where"].values()))]
+        }
 
 
 # ---------------------------------------------------------------------------