From aa3c45f1e76b00abb1f23b0a6c4ed0fcdd480459 Mon Sep 17 00:00:00 2001 From: Vineeth Sai Date: Tue, 8 Sep 2026 12:54:38 -0700 Subject: [PATCH] test(proxy): pass the new _reset_expired_window arguments 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. --- .../proxy/common_utils/test_reset_budget_job.py | 16 ++++++++++++++-- 1 file changed, 14 insertions(+), 2 deletions(-) diff --git a/tests/test_litellm/proxy/common_utils/test_reset_budget_job.py b/tests/test_litellm/proxy/common_utils/test_reset_budget_job.py index a91a2d1a4c7..0ec297c4505 100644 --- a/tests/test_litellm/proxy/common_utils/test_reset_budget_job.py +++ b/tests/test_litellm/proxy/common_utils/test_reset_budget_job.py @@ -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) ]