scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-07 08:26:10 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 36

test(proxy): cover default-tier end-user counter reset with rollover

Browse source
This commit is contained in:
mateo-berri 2026-09-04 16:42:29 -07:00
parent 9c8594c7b8
commit 976f8625f3
1 changed files with 32 additions and 0 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

32
tests/test_litellm/proxy/common_utils/test_reset_budget_job.py
View file

@ -3054,6 +3054,38 @@ def test_budget_cascade_carries_enduser_overage_when_rollover_enabled(
} in enduser_writes
def test_budget_cascade_carries_default_tier_enduser_counter_when_rollover_enabled(
rollover_enabled, reset_budget_job, mock_prisma_client, monkeypatch
):
"""An end user on the default budget (no budget_id on its row) 5 over the cap
keeps a counter of 5 in the next window and loses its cached object."""
import litellm
counter_cache: Final = _make_counter_invalidation_job(monkeypatch)
monkeypatch.setattr(litellm, "max_end_user_budget_id", "default-enduser-budget")
mock_prisma_client.data["budget"] = [
_budget_row(budget_id="default-enduser-budget", budget_duration="1d", max_budget=10.0)
]
implicit_enduser: Final = type(
"EndUserRow",
(),
{
"spend": 15.0,
"user_id": "enduser-implicit",
"budget_id": None,
"model_dump": lambda self=None: {"spend": 15.0, "user_id": "enduser-implicit", "budget_id": None, "blocked": False},
},
)
mock_prisma_client.db.litellm_endusertable.set_find_many_results([implicit_enduser])
asyncio.run(reset_budget_job.reset_budget_for_litellm_budget_table())
counter_cache.in_memory_cache.set_cache.assert_any_call(key="spend:end_user:enduser-implicit", value=5.0, ttl=60)
counter_cache.redis_cache.async_set_cache.assert_any_await(key="spend:end_user:enduser-implicit", value=5.0, ttl=60)
deleted: Final = {call.kwargs.get("key") for call in counter_cache.user_api_key_cache.async_delete_cache.await_args_list}
assert "end_user_id:enduser-implicit" in deleted
def _replay_spend_writes(writes, spend):
"""Apply the queued update_many statements in order, the way the DB
transaction executes them, and return the row's final spend."""