scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-16 23:41:43 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 102

Merge pull request #41359 from BerriAI/litellm_fix_budget_reset_decrement_tests

Browse source 
test(proxy): assert budget resets decrement the cleared spend
This commit is contained in:
yuneng-jiang 2026-09-15 20:27:35 -07:00 • committed by GitHub
commit 9d54c2e159
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
1 changed files with 20 additions and 4 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

24
tests/litellm_utils_tests/test_proxy_budget_reset.py
View file

@ -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()))]
}
# ---------------------------------------------------------------------------