From 53478e86517b8ec4ea5801feb76b4468243ecec1 Mon Sep 17 00:00:00 2001 From: Emerson Gomes Date: Tue, 15 Sep 2026 13:39:49 -0500 Subject: [PATCH] test: align budget hotpath checks with awaited Redis flush --- litellm/router_strategy/budget_limiter.py | 4 ++-- .../test_budget_limiter_hotpath.py | 22 +++++++++++-------- 2 files changed, 15 insertions(+), 11 deletions(-) diff --git a/litellm/router_strategy/budget_limiter.py b/litellm/router_strategy/budget_limiter.py index 43ee2658a30..089e95d10e3 100644 --- a/litellm/router_strategy/budget_limiter.py +++ b/litellm/router_strategy/budget_limiter.py @@ -460,8 +460,8 @@ class RouterBudgetLimiting(CustomLogger): ) try: await redis_cache.async_increment_pipeline(increment_list=increment_list) - except Exception: - verbose_router_logger.exception("Error pushing queued Redis increment operations to Redis") + except Exception as error: + log_redis_failure(verbose_router_logger, logging.ERROR, "Error syncing in-memory cache with Redis", error) await self._requeue_detached_increment_operations() return False await self._clear_detached_increment_operations() diff --git a/tests/test_litellm/router_strategy/test_budget_limiter_hotpath.py b/tests/test_litellm/router_strategy/test_budget_limiter_hotpath.py index 4cc8fe78811..272ecd4179a 100644 --- a/tests/test_litellm/router_strategy/test_budget_limiter_hotpath.py +++ b/tests/test_litellm/router_strategy/test_budget_limiter_hotpath.py @@ -338,7 +338,9 @@ async def test_sync_refused_by_the_open_circuit_breaker_is_quiet_and_leaks_no_ta assert caplog.records == [] unretrieved.assert_not_called() - assert limiter.redis_increment_operation_queue == [] + assert limiter.redis_increment_operation_queue == [ + {"key": "provider_spend:openai:1d", "increment_value": 0.5, "ttl": 60} + ] assert redis_cache.async_increment_pipeline.await_count == 1 @@ -353,25 +355,27 @@ async def _limiter_with_redis(redis_cache: MagicMock) -> RouterBudgetLimiting: @pytest.mark.asyncio -async def test_push_returns_before_redis_answers(disable_budget_sync): - """The push runs inside the request success callback, so it must hand the Redis round trip to a task instead of waiting on it.""" +async def test_push_waits_for_redis_before_completing(disable_budget_sync): + redis_started = asyncio.Event() redis_answered = asyncio.Event() async def wait_for_redis(**_: object) -> None: + redis_started.set() await redis_answered.wait() redis_cache = MagicMock(spec=RedisCache) redis_cache.async_increment_pipeline = AsyncMock(side_effect=wait_for_redis) limiter = await _limiter_with_redis(redis_cache) - await asyncio.wait_for(limiter._push_in_memory_increments_to_redis(), timeout=1) - await asyncio.sleep(0) - - assert not redis_answered.is_set() + push_task = asyncio.create_task(limiter._push_in_memory_increments_to_redis()) + await asyncio.wait_for(redis_started.wait(), timeout=1) + assert not push_task.done() + assert limiter._detached_increment_operations is not None + redis_answered.set() + assert await asyncio.wait_for(push_task, timeout=1) is True assert redis_cache.async_increment_pipeline.await_count == 1 assert limiter.redis_increment_operation_queue == [] - redis_answered.set() - await asyncio.gather(*(task for task in asyncio.all_tasks() if task is not asyncio.current_task())) + assert limiter._detached_increment_operations is None @pytest.mark.asyncio