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 103

test(proxy): pass the new _reset_expired_window arguments

Browse source 
Upstream added prisma_client, entity_type and entity_id so the reset can
roll the window's LiteLLM_BudgetWindowSpend row. The two call sites in
TestBudgetWindowResetAtOffsets still used the old signature and raised
TypeError. The row roll is best effort inside its own try, so a MagicMock
client exercises the reset decision without a database.

tests/test_litellm/proxy/common_utils/test_reset_budget_job.py: 140 passed.
This commit is contained in:
Vineeth Sai 2026-09-08 12:54:38 -07:00
parent 7e9e93521b
commit aa3c45f1e7
1 changed files with 14 additions and 2 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

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

@ -13,7 +13,7 @@ import pytest
from litellm.caching.dual_cache import DualCache
from litellm.proxy._types import LiteLLM_VerificationToken
from litellm.proxy._types import Litellm_EntityType, LiteLLM_VerificationToken
from litellm.proxy.common_utils import reset_budget_job as reset_budget_job_module
from litellm.constants import (
PROXY_BUDGET_RESCHEDULER_MIN_TIME,
@ -3211,12 +3211,17 @@ class TestBudgetWindowResetAtOffsets:
@staticmethod
async def _expired(window: dict, now: datetime, tz: str = "UTC") -> bool:
# The spend-row roll is best effort and wrapped in its own try, so a mock
# client here exercises the reset decision without touching a database.
return await ResetBudgetJob._reset_expired_window(
window,
"spend:key:tok-1:window:" + window["budget_duration"],
DualCache(),
now,
BudgetResetSettings(timezone=tz, reset_time_of_day=dt_time(0, 0)),
MagicMock(),
Litellm_EntityType.KEY,
"tok-1",
)
@pytest.mark.asyncio
@ -3283,7 +3288,14 @@ class TestBudgetWindowResetAtOffsets:
results = [
await ResetBudgetJob._reset_expired_window(
window, f"spend:key:tok-1:window:{window['budget_duration']}", cache, now, settings
window,
f"spend:key:tok-1:window:{window['budget_duration']}",
cache,
now,
settings,
MagicMock(),
Litellm_EntityType.KEY,
"tok-1",
)
for window in (due, not_due)
]