diff --git a/tests/test_litellm/caching/test_redis_cache.py b/tests/test_litellm/caching/test_redis_cache.py index c1e3240adb7..e811399f901 100644 --- a/tests/test_litellm/caching/test_redis_cache.py +++ b/tests/test_litellm/caching/test_redis_cache.py @@ -77,6 +77,52 @@ async def test_async_delete_cache_applies_namespace( mock_redis_instance.delete.assert_awaited_once_with(expected_key) +@pytest.mark.asyncio +async def test_async_reset_preserving_delta_evals_with_namespaced_key_and_string_args( + monkeypatch, redis_no_ping +): + """The GET/compute/SET has to run as one Lua call, not separate round trips, or a + concurrent async_increment between them would be exactly the race this method exists + to close. Namespacing and str-ifying every ARGV also has to happen, or redis-py's own + encoding (or an ACL scoped to the namespace prefix) breaks the call outright.""" + monkeypatch.setenv("REDIS_HOST", "https://my-test-host") + redis_cache = RedisCache(namespace="litellm") + mock_redis_instance = AsyncMock() + mock_redis_instance.eval.return_value = "12.5" + + with patch.object( + redis_cache, "init_async_client", return_value=mock_redis_instance + ): + result = await redis_cache.async_reset_preserving_delta( + key="spend:key:abc", new_base=10.0, snapshot=7.5, ttl=60 + ) + + assert result == 12.5 + mock_redis_instance.eval.assert_awaited_once() + lua, numkeys, key, new_base_arg, snapshot_arg, ttl_arg = mock_redis_instance.eval.await_args.args + assert numkeys == 1 + assert key == "litellm:spend:key:abc" + assert (new_base_arg, snapshot_arg, ttl_arg) == ("10.0", "7.5", "60") + assert "GET" in lua and "SET" in lua and "EXPIRE" in lua + + +@pytest.mark.asyncio +async def test_async_reset_preserving_delta_decodes_a_bytes_result(monkeypatch, redis_no_ping): + """redis-py returns EVAL results as bytes unless decode_responses is set; a caller that + compares the return value to a float must not have to know that.""" + monkeypatch.setenv("REDIS_HOST", "https://my-test-host") + redis_cache = RedisCache() + mock_redis_instance = AsyncMock() + mock_redis_instance.eval.return_value = b"10.0" + + with patch.object( + redis_cache, "init_async_client", return_value=mock_redis_instance + ): + result = await redis_cache.async_reset_preserving_delta(key="k", new_base=10.0, snapshot=10.0, ttl=60) + + assert result == 10.0 + + @pytest.mark.parametrize("namespace", [None, "litellm"]) def test_delete_cache_applies_namespace(namespace, monkeypatch, redis_no_ping): """delete_cache must prefix keys with the namespace, matching every other 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 15cf6c5fd87..a974000bbf8 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 @@ -3479,6 +3479,30 @@ def test_invalidate_spend_counter_retries_then_deletes_on_persistent_redis_failu spend_counter_cache.redis_cache.async_delete_cache.assert_awaited_once_with(key="spend:key:sk-inflated") +def test_invalidate_spend_counter_swallows_a_delete_failure_after_reset_already_failed(monkeypatch): + """The fallback delete is itself best-effort: if it also fails (e.g. the same + outage that broke the reset), there is nothing left to try, and the failure + must be logged and swallowed rather than propagate out of the reset job.""" + spend_counter_cache = MagicMock() + spend_counter_cache.in_memory_cache.set_cache = MagicMock() + spend_counter_cache.redis_cache = MagicMock() + spend_counter_cache.redis_cache.async_get_cache = AsyncMock(return_value=80.0) + spend_counter_cache.redis_cache.async_reset_preserving_delta = AsyncMock( + side_effect=RuntimeError("elasticache timeout") + ) + spend_counter_cache.redis_cache.async_delete_cache = AsyncMock(side_effect=RuntimeError("elasticache timeout")) + + fake_module = types.ModuleType("litellm.proxy.proxy_server") + fake_module.spend_counter_cache = spend_counter_cache + monkeypatch.setitem(sys.modules, "litellm.proxy.proxy_server", fake_module) + monkeypatch.setattr(asyncio, "sleep", AsyncMock()) + + # must not raise + asyncio.run(ResetBudgetJob._invalidate_spend_counter("spend:key:sk-double-failure", new_spend=0.0)) + + spend_counter_cache.redis_cache.async_delete_cache.assert_awaited_once_with(key="spend:key:sk-double-failure") + + def test_invalidate_spend_counter_recovers_after_a_transient_redis_failure(monkeypatch): """A reset that fails once and then succeeds must not fall back to delete: the counter ends up reset to the real value, not merely absent.""" diff --git a/tests/test_litellm/proxy/test_budget_reservation.py b/tests/test_litellm/proxy/test_budget_reservation.py index a3f33c2b07e..64756558e39 100644 --- a/tests/test_litellm/proxy/test_budget_reservation.py +++ b/tests/test_litellm/proxy/test_budget_reservation.py @@ -2788,6 +2788,49 @@ async def test_release_budget_reservation_on_cancel_swallows_release_errors(): await release_budget_reservation_on_cancel(reservation) +@pytest.mark.asyncio +async def test_release_budget_reservation_on_cancel_swallows_a_second_cancellation_while_shielded(): + # A second CancelledError arriving while the shielded reconcile is in flight must not + # propagate: the reconcile keeps running detached regardless, and there is nothing more + # for this call to do but return. + reservation = { + "reserved_cost": 3.0, + "entries": [{"counter_key": "spend:key:key-cancel-twice"}], + "finalized": False, + "input_cost": 0.5, + } + with patch( + "litellm.proxy.spend_tracking.budget_reservation.reconcile_budget_reservation", + new=AsyncMock(side_effect=asyncio.CancelledError()), + ): + # must return without raising + await release_budget_reservation_on_cancel(reservation) + + +@pytest.mark.asyncio +async def test_release_budget_reservation_on_cancel_swallows_invalidate_failure_after_every_retry_fails(): + # If both the reconcile retries and the invalidate fallback fail (e.g. a persistent + # outage), there is nothing left to try: the failure must be logged and swallowed, + # not propagated, and the reservation still ends up finalized so it is not reprocessed. + reservation = { + "reserved_cost": 3.0, + "entries": [{"counter_key": "spend:key:key-cancel-double-failure"}], + "finalized": False, + "input_cost": 0.5, + } + with patch( + "litellm.proxy.spend_tracking.budget_reservation.reconcile_budget_reservation", + new=AsyncMock(side_effect=RuntimeError("redis down")), + ), patch( + "litellm.proxy.spend_tracking.budget_reservation.invalidate_budget_reservation_counters", + new=AsyncMock(side_effect=RuntimeError("redis still down")), + ): + # must return without raising + await release_budget_reservation_on_cancel(reservation) + + assert reservation["finalized"] is True + + @pytest.mark.asyncio async def test_release_budget_reservation_on_cancel_invalidates_counter_when_reconcile_persistently_fails( spend_counter_state,