litellm/tests/unit/proxy/proxy_server/test_spend_counters.py

"""Behavior pins for spend-counter helpers in proxy_server.

Pins covered:
- ``get_current_spend``
- ``increment_spend_counters``
- ``_reconcile_budget_reservation_for_counter_update``
- ``_prepare_end_user_and_tag_spend_increments``
- ``_prepare_org_spend_increment``
- ``_prepare_unreserved_spend_counter_increment``
- ``_prepare_spend_counter_increment``
- ``_prepare_window_spend_counter_increment``
- ``_apply_spend_counter_increments``
- ``_ensure_spend_counter_initialized``
- ``_get_source_cache_base_spend``
- ``_ensure_window_spend_counter_initialized``
- ``_is_spend_counter_cache_warm``
- ``_increment_spend_counter_cache``
- ``_invalidate_spend_counter``
- ``update_cache``
"""

from __future__ import annotations

import asyncio
from collections.abc import Awaitable, Callable
from datetime import datetime
from typing import Final, Literal, cast
from unittest.mock import AsyncMock, MagicMock, call

import pytest
from fakeredis import FakeAsyncRedis, FakeServer
from opentelemetry.sdk.trace import TracerProvider
from opentelemetry.sdk.trace.export import SimpleSpanProcessor
from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter
from opentelemetry.trace import StatusCode

import litellm
import litellm.proxy.proxy_server as ps
from litellm._service_logger import ServiceLogging
from litellm._internal_context import current_service_target
from litellm.caching.in_memory_cache import InMemoryCache
from litellm.caching.redis_cache import RedisCache, RedisCircuitBreaker
from litellm.constants import GLOBAL_PROXY_SPEND_CACHE_KEY
from litellm.integrations.opentelemetry import OpenTelemetry, OpenTelemetryConfig
from litellm.models.end_user import LiteLLM_EndUserTable
from litellm.models.tag import LiteLLM_TagTable
from litellm.models.team import LiteLLM_TeamTable, LiteLLM_TeamTableCachedObj
from litellm.models.user import LiteLLM_UserTable
from litellm.proxy.common_utils.auth_cache_invalidation_pubsub import evict_local
from litellm.proxy.common_utils.cache_pydantic_utils import CacheCodec
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache, end_user_cache_key, tag_cache_key
from litellm.proxy.management_endpoints.team_endpoints import refresh_cached_team

from .conftest import normalize


def _make_spend_counter_cache(
    *,
    redis_get_value=None,
    redis_get_side_effect=None,
    redis_increment_value=None,
    redis_increment_side_effect=None,
    in_memory_value=None,
    with_redis: bool = True,
):
    cache = MagicMock()
    cache.in_memory_cache = MagicMock()
    cache.in_memory_cache.get_cache = MagicMock(return_value=in_memory_value)
    cache.in_memory_cache.set_cache = MagicMock()
    cache.in_memory_cache.delete_cache = MagicMock()
    if with_redis:
        cache.redis_cache = MagicMock()
        cache.redis_cache.async_get_cache = AsyncMock(return_value=redis_get_value, side_effect=redis_get_side_effect)
        cache.redis_cache.async_increment = AsyncMock(
            return_value=redis_increment_value,
            side_effect=redis_increment_side_effect,
        )
        cache.redis_cache.async_delete_cache = AsyncMock()
        cache.redis_cache.async_set_cache = AsyncMock()
        cache.redis_cache.async_set_max = AsyncMock()
        cache.redis_cache.async_increment_pipeline = AsyncMock(return_value=None)
        cache.redis_cache.get_ttl = MagicMock(return_value=None)
    else:
        cache.redis_cache = None
    cache.async_increment_cache = AsyncMock(return_value=redis_increment_value)
    cache.async_get_cache = AsyncMock(return_value=None)
    cache.async_set_cache = AsyncMock()
    cache.async_delete_cache = AsyncMock()
    cache.async_set_cache_pipeline = AsyncMock()
    return cache


def _make_user_api_key_cache(get_value=None, get_side_effect=None):
    cache = MagicMock()
    cache.async_get_cache = AsyncMock(return_value=get_value, side_effect=get_side_effect)
    cache.async_batch_get_cache = AsyncMock(side_effect=lambda keys, **_: [get_value for _ in keys])
    cache.async_set_cache_pipeline = AsyncMock()
    return cache


# ---------------------------------------------------------------------------
# get_current_spend
# ---------------------------------------------------------------------------


@pytest.mark.asyncio
async def test_get_current_spend_reads_redis_first(monkeypatch):
    fake_cache = _make_spend_counter_cache(redis_get_value=42.5)
    monkeypatch.setattr(ps, "spend_counter_cache", fake_cache)

    result = await ps.get_current_spend(counter_key="spend:key:abc", fallback_spend=0.0)

    observed = {
        "value": result,
        "redis_called": fake_cache.redis_cache.async_get_cache.called,
        "in_memory_called": fake_cache.in_memory_cache.get_cache.called,
    }
    assert normalize(observed) == {
        "value": 42.5,
        "redis_called": True,
        "in_memory_called": False,
    }


@pytest.mark.asyncio
async def test_get_current_spend_redis_error_falls_back_to_in_memory(monkeypatch):
    fake_cache = _make_spend_counter_cache(
        redis_get_side_effect=RuntimeError("redis down"),
        in_memory_value=17.0,
    )
    monkeypatch.setattr(ps, "spend_counter_cache", fake_cache)

    result = await ps.get_current_spend(counter_key="spend:key:abc", fallback_spend=99.0)
    assert result == 17.0


@pytest.mark.asyncio
async def test_get_current_spend_floors_stale_low_counter_against_db(monkeypatch):
    """A Redis counter left stale-low by a Redis restart must not admit a key
    whose authoritative DB spend is already over budget. With max_budget set,
    get_current_spend re-checks the DB and returns the higher recorded spend."""
    fake_cache = _make_spend_counter_cache(redis_get_value=2.0)
    monkeypatch.setattr(ps, "spend_counter_cache", fake_cache)
    from_db = AsyncMock(return_value=12.0)
    monkeypatch.setattr(ps.SpendCounterReseed, "from_db", from_db)

    result = await ps.get_current_spend(
        counter_key="spend:key:abc",
        fallback_spend=12.0,
        max_budget=10.0,
    )

    assert result == 12.0
    assert from_db.await_count == 1
    # the stale counter is repaired up to the authoritative DB value via a
    # monotonic set-max so other workers read the corrected total, and a
    # concurrent increment cannot be clobbered
    fake_cache.redis_cache.async_set_max.assert_awaited_once_with(key="spend:key:abc", value=12.0)


@pytest.mark.asyncio
async def test_get_current_spend_no_db_recheck_when_counter_healthy(monkeypatch):
    """A healthy counter (at or above the caller's recorded spend) is trusted
    without a DB read, so under-budget traffic stays off the DB path."""
    fake_cache = _make_spend_counter_cache(redis_get_value=5.0)
    monkeypatch.setattr(ps, "spend_counter_cache", fake_cache)
    from_db = AsyncMock(return_value=99.0)
    monkeypatch.setattr(ps.SpendCounterReseed, "from_db", from_db)

    result = await ps.get_current_spend(
        counter_key="spend:key:abc",
        fallback_spend=3.0,
        max_budget=10.0,
    )

    assert result == 5.0
    assert from_db.await_count == 0


@pytest.mark.asyncio
async def test_get_current_spend_no_floor_without_max_budget(monkeypatch):
    """Without max_budget the read-time DB floor is skipped: callers that only
    read spend (alerts, soft budgets) keep the cheap counter-only behavior."""
    fake_cache = _make_spend_counter_cache(redis_get_value=2.0)
    monkeypatch.setattr(ps, "spend_counter_cache", fake_cache)
    from_db = AsyncMock(return_value=12.0)
    monkeypatch.setattr(ps.SpendCounterReseed, "from_db", from_db)

    result = await ps.get_current_spend(counter_key="spend:key:abc", fallback_spend=12.0)

    assert result == 2.0
    assert from_db.await_count == 0


@pytest.mark.asyncio
async def test_get_current_spend_floor_admits_after_reset(monkeypatch):
    """Right after a weekly reset the counter is 0 while the per-worker cached
    spend can still be last week's value. The DB floor reads the reset spend (0)
    and admits, so reset keys are not over-blocked."""
    fake_cache = _make_spend_counter_cache(redis_get_value=0.0)
    monkeypatch.setattr(ps, "spend_counter_cache", fake_cache)
    from_db = AsyncMock(return_value=0.0)
    monkeypatch.setattr(ps.SpendCounterReseed, "from_db", from_db)

    result = await ps.get_current_spend(
        counter_key="spend:key:abc",
        fallback_spend=12.0,
        max_budget=10.0,
    )

    assert result == 0.0
    assert from_db.await_count == 1
    # counter already matches the DB (reset to 0); nothing to repair, so no write
    fake_cache.redis_cache.async_set_max.assert_not_called()


@pytest.mark.asyncio
async def test_get_current_spend_floor_caches_db_read(monkeypatch):
    """A persistently stale-low counter must not drive a DB read per request:
    the authoritative spend is cached in-process and reused within the window."""
    cache = ps.DualCache()
    cache.redis_cache = MagicMock()
    cache.redis_cache.async_get_cache = AsyncMock(return_value=2.0)
    monkeypatch.setattr(ps, "spend_counter_cache", cache)
    from_db = AsyncMock(return_value=12.0)
    monkeypatch.setattr(ps.SpendCounterReseed, "from_db", from_db)

    first = await ps.get_current_spend(counter_key="spend:key:abc", fallback_spend=12.0, max_budget=10.0)
    second = await ps.get_current_spend(counter_key="spend:key:abc", fallback_spend=12.0, max_budget=10.0)

    assert first == 12.0
    assert second == 12.0
    assert from_db.await_count == 1


@pytest.mark.asyncio
@pytest.mark.parametrize("counter_key", ("spend:end_user:e1", "spend:tag:t1"))
async def test_get_current_spend_floors_end_user_tag_against_fallback(monkeypatch, counter_key):
    """Tag counters have no DB row (from_db returns None), and an end-user counter has
    none to read without a DB client. When such a counter is stale-low, enforcement
    falls back to the caller's recorded spend (loaded fresh in auth) instead of
    trusting the stale counter."""
    fake_cache: Final = _make_spend_counter_cache(redis_get_value=2.0)
    monkeypatch.setattr(ps, "spend_counter_cache", fake_cache)
    monkeypatch.setattr(ps, "prisma_client", None)
    monkeypatch.setattr(ps.SpendCounterReseed, "from_db", AsyncMock(return_value=None))

    result: Final = await ps.get_current_spend(
        counter_key=counter_key,
        fallback_spend=20.0,
        max_budget=10.0,
    )

    assert result == 20.0
    # no DB row to repair against, so the shared counter is left untouched
    fake_cache.redis_cache.async_set_max.assert_not_called()


def _make_prisma_with_end_user_row(spend: float | None):
    prisma: Final = MagicMock()
    prisma.db.litellm_endusertable.find_unique = AsyncMock(
        return_value=None if spend is None else MagicMock(spend=spend)
    )
    return prisma


@pytest.mark.asyncio
async def test_get_current_spend_end_user_floor_admits_after_a_reset_on_a_stale_worker(monkeypatch):
    """The reset job zeroes LiteLLM_EndUserTable.spend and the shared counter, but it
    evicts the cached end-user object only on the worker that ran the reset. Every
    other worker still passes the pre-reset spend as fallback_spend, and that stale
    copy must not out-vote the reset row."""
    fake_cache: Final = _make_spend_counter_cache(redis_get_value=0.0)
    monkeypatch.setattr(ps, "spend_counter_cache", fake_cache)
    prisma: Final = _make_prisma_with_end_user_row(spend=0.0)
    monkeypatch.setattr(ps, "prisma_client", prisma)

    result = await ps.get_current_spend(
        counter_key="spend:end_user:customer-42",
        fallback_spend=0.000032,
        max_budget=0.00003,
        fallback_authoritative=True,
    )

    assert result == 0.0
    prisma.db.litellm_endusertable.find_unique.assert_awaited_once_with(where={"user_id": "customer-42"})
    fake_cache.redis_cache.async_set_max.assert_not_called()


@pytest.mark.asyncio
async def test_get_current_spend_end_user_floor_repairs_a_stale_low_counter(monkeypatch):
    """After a Redis restart the end-user counter can sit below the recorded spend;
    the row wins and the shared counter is raised so other workers stop admitting on
    the stale value."""
    fake_cache: Final = _make_spend_counter_cache(redis_get_value=2.0)
    monkeypatch.setattr(ps, "spend_counter_cache", fake_cache)
    monkeypatch.setattr(ps, "prisma_client", _make_prisma_with_end_user_row(spend=12.0))

    result: Final = await ps.get_current_spend(
        counter_key="spend:end_user:customer-42",
        fallback_spend=12.0,
        max_budget=10.0,
    )

    assert result == 12.0
    fake_cache.redis_cache.async_set_max.assert_awaited_once_with(key="spend:end_user:customer-42", value=12.0)


@pytest.mark.asyncio
async def test_get_current_spend_end_user_without_a_row_keeps_the_cached_spend(monkeypatch):
    fake_cache: Final = _make_spend_counter_cache(redis_get_value=0.0)
    monkeypatch.setattr(ps, "spend_counter_cache", fake_cache)
    monkeypatch.setattr(ps, "prisma_client", _make_prisma_with_end_user_row(spend=None))

    result: Final = await ps.get_current_spend(
        counter_key="spend:end_user:customer-42",
        fallback_spend=20.0,
        max_budget=10.0,
    )

    assert result == 20.0
    fake_cache.redis_cache.async_set_max.assert_not_called()


@pytest.mark.asyncio
async def test_get_current_spend_floors_window_against_spend_logs(monkeypatch):
    """Per-window counters have no DB row but aggregate from spend logs. A
    stale-low window counter is floored to (and repaired up to) the logged
    window spend, even though the caller's fallback is 0."""
    from datetime import datetime, timezone

    fake_cache = _make_spend_counter_cache(redis_get_value=2.0)
    monkeypatch.setattr(ps, "spend_counter_cache", fake_cache)
    monkeypatch.setattr(ps.SpendCounterReseed, "from_db", AsyncMock(return_value=None))
    wfsl = AsyncMock(return_value=15.0)
    monkeypatch.setattr(ps.SpendCounterReseed, "window_from_spend_logs", wfsl)

    counter_key = "spend:key:tok:window:7d"
    result = await ps.get_current_spend(
        counter_key=counter_key,
        fallback_spend=0.0,
        max_budget=10.0,
        window_entity_type="Key",
        window_entity_id="tok",
        window_start=datetime(2026, 1, 1, tzinfo=timezone.utc),
    )

    assert result == 15.0
    assert wfsl.await_count == 1
    fake_cache.redis_cache.async_set_max.assert_awaited_once_with(key=counter_key, value=15.0)


def _make_window_spend_prisma(row=None, spend_logs_total=0.0):
    prisma = MagicMock()
    prisma.db.litellm_budgetwindowspend.find_unique = AsyncMock(return_value=row)
    prisma.db.litellm_spendlogs.group_by = AsyncMock(
        return_value=[{"api_key": "tok", "_sum": {"spend": spend_logs_total}}]
    )
    return prisma


@pytest.mark.asyncio
async def test_get_current_spend_floors_window_against_maintained_row(monkeypatch):
    """The floor re-check runs every few seconds per pod, so the window branch
    must read the maintained row and leave the unindexed spend-logs scan alone."""
    from datetime import timezone
    from types import SimpleNamespace

    window_start = datetime(2026, 1, 1, tzinfo=timezone.utc)
    fake_prisma = _make_window_spend_prisma(
        row=SimpleNamespace(window_start=window_start, spend=15.0),
        spend_logs_total=100.0,
    )
    fake_cache = _make_spend_counter_cache(redis_get_value=2.0)
    monkeypatch.setattr(ps, "spend_counter_cache", fake_cache)
    monkeypatch.setattr(ps, "prisma_client", fake_prisma)

    counter_key = "spend:key:tok:window:7d"
    result = await ps.get_current_spend(
        counter_key=counter_key,
        fallback_spend=0.0,
        max_budget=10.0,
        window_entity_type="Key",
        window_entity_id="tok",
        window_duration="7d",
        window_start=window_start,
    )

    assert result == 15.0
    fake_prisma.db.litellm_spendlogs.group_by.assert_not_awaited()
    fake_cache.redis_cache.async_set_max.assert_awaited_once_with(key=counter_key, value=15.0)


@pytest.mark.asyncio
async def test_get_current_spend_floors_window_against_logs_when_row_stale(monkeypatch):
    """A row left behind at a crossed window boundary must not be read as the
    current window's spend; the aggregate stays the fallback."""
    from datetime import timedelta, timezone
    from types import SimpleNamespace

    window_start = datetime(2026, 1, 8, tzinfo=timezone.utc)
    fake_prisma = _make_window_spend_prisma(
        row=SimpleNamespace(window_start=window_start - timedelta(days=7), spend=999.0),
        spend_logs_total=15.0,
    )
    fake_cache = _make_spend_counter_cache(redis_get_value=2.0)
    monkeypatch.setattr(ps, "spend_counter_cache", fake_cache)
    monkeypatch.setattr(ps, "prisma_client", fake_prisma)

    result = await ps.get_current_spend(
        counter_key="spend:key:tok:window:7d",
        fallback_spend=0.0,
        max_budget=10.0,
        window_entity_type="Key",
        window_entity_id="tok",
        window_duration="7d",
        window_start=window_start,
    )

    assert result == 15.0
    fake_prisma.db.litellm_spendlogs.group_by.assert_awaited_once()


@pytest.mark.asyncio
async def test_get_current_spend_fail_closed_rejects_when_unverifiable(monkeypatch):
    """With fail_closed_budget_enforcement on, an admit decision backed only by a
    per-pod fallback (Redis unreachable and DB unreadable) is rejected with 503
    rather than admitted on an unverifiable budget."""
    from fastapi import HTTPException

    fake_cache = _make_spend_counter_cache(redis_get_side_effect=RuntimeError("redis down"))
    monkeypatch.setattr(ps, "spend_counter_cache", fake_cache)
    monkeypatch.setattr(ps, "general_settings", {"fail_closed_budget_enforcement": True})
    monkeypatch.setattr(ps.SpendCounterReseed, "coalesced", AsyncMock(return_value=None))

    with pytest.raises(HTTPException) as exc:
        await ps.get_current_spend(counter_key="spend:key:abc", fallback_spend=1.0, max_budget=10.0)
    assert exc.value.status_code == 503


@pytest.mark.asyncio
async def test_get_current_spend_fail_closed_off_admits_when_unverifiable(monkeypatch):
    """Default (flag off): an unverifiable read keeps the existing behavior and
    admits using the cached fallback — no new rejection."""
    fake_cache = _make_spend_counter_cache(redis_get_side_effect=RuntimeError("redis down"))
    monkeypatch.setattr(ps, "spend_counter_cache", fake_cache)
    monkeypatch.setattr(ps, "general_settings", {})
    monkeypatch.setattr(ps.SpendCounterReseed, "coalesced", AsyncMock(return_value=None))

    result = await ps.get_current_spend(counter_key="spend:key:abc", fallback_spend=1.0, max_budget=10.0)
    assert result == 1.0


@pytest.mark.asyncio
async def test_get_current_spend_fail_closed_admits_when_redis_verified(monkeypatch):
    """Fail-closed only rejects unverifiable reads: a value served by Redis is
    authoritative, so an under-budget request is admitted normally."""
    fake_cache = _make_spend_counter_cache(redis_get_value=1.0)
    monkeypatch.setattr(ps, "spend_counter_cache", fake_cache)
    monkeypatch.setattr(ps, "general_settings", {"fail_closed_budget_enforcement": True})

    result = await ps.get_current_spend(counter_key="spend:key:abc", fallback_spend=1.0, max_budget=10.0)
    assert result == 1.0


@pytest.mark.asyncio
async def test_get_current_spend_fail_closed_allows_authoritative_fallback(monkeypatch):
    """End-user/tag callers pass fallback_authoritative=True (their spend is
    loaded fresh from the DB in auth), so fail-closed does not reject them even
    when the counter path is unreadable."""
    fake_cache = _make_spend_counter_cache(redis_get_side_effect=RuntimeError("redis down"))
    monkeypatch.setattr(ps, "spend_counter_cache", fake_cache)
    monkeypatch.setattr(ps, "general_settings", {"fail_closed_budget_enforcement": True})
    monkeypatch.setattr(ps.SpendCounterReseed, "coalesced", AsyncMock(return_value=None))

    result = await ps.get_current_spend(
        counter_key="spend:end_user:e1",
        fallback_spend=1.0,
        max_budget=10.0,
        fallback_authoritative=True,
    )
    assert result == 1.0


@pytest.mark.asyncio
async def test_get_current_spend_strict_floors_when_fallback_also_stale(monkeypatch):
    """Strict mode closes the both-stale gap: when the counter AND the caller's
    cached spend are both stale-low (cheap guard would skip), strict mode still
    re-checks the authoritative DB and enforces against it."""
    fake_cache = _make_spend_counter_cache(redis_get_value=0.00001)
    monkeypatch.setattr(ps, "spend_counter_cache", fake_cache)
    monkeypatch.setattr(ps, "general_settings", {"fail_closed_budget_enforcement": True})
    from_db = AsyncMock(return_value=0.5)
    monkeypatch.setattr(ps.SpendCounterReseed, "from_db", from_db)

    # fallback == current, so the default cheap guard would NOT re-check
    result = await ps.get_current_spend(
        counter_key="spend:team:t1",
        fallback_spend=0.00001,
        max_budget=0.0002,
    )

    assert result == 0.5
    assert from_db.await_count == 1


# ---------------------------------------------------------------------------
# increment_spend_counters
# ---------------------------------------------------------------------------


@pytest.mark.asyncio
async def test_increment_spend_counters_increments_all_buckets(monkeypatch):
    fake_cache = _make_spend_counter_cache(redis_get_value=None, redis_increment_value=5.0)
    fake_user_cache = _make_user_api_key_cache(get_value=None)
    monkeypatch.setattr(ps, "spend_counter_cache", fake_cache)
    monkeypatch.setattr(ps, "user_api_key_cache", fake_user_cache)
    monkeypatch.setattr(ps, "prisma_client", None)

    async def _fake_coalesced(**kwargs):
        return None

    monkeypatch.setattr(ps.SpendCounterReseed, "coalesced", AsyncMock(side_effect=_fake_coalesced))

    await ps.increment_spend_counters(
        token="hashed-tok",
        team_id="t1",
        user_id="u1",
        response_cost=5.0,
    )

    pipeline = fake_cache.redis_cache.async_increment_pipeline
    pipeline.assert_awaited_once()
    increment_list = pipeline.await_args.kwargs["increment_list"]
    assert {op["key"] for op in increment_list} == {
        "spend:key:hashed-tok",
        "spend:team:t1",
        "spend:team_member:u1:t1",
        "spend:user:u1",
    }
    assert all(op["increment_value"] == 5.0 for op in increment_list)
    observed = {
        "redis_increment_called": fake_cache.redis_cache.async_increment.called,
        "pipeline_calls": pipeline.await_count,
        "user_cache_used": fake_user_cache.async_get_cache.called,
    }
    assert normalize(observed) == {
        "redis_increment_called": False,
        "pipeline_calls": 1,
        "user_cache_used": True,
    }


class _ConcurrencyProbe:
    """Stand-in for redis_cache.async_get_cache that pins concurrency.

    Each warm-check read registers itself as in-flight and blocks on ``release``
    until the test lets it proceed. ``all_arrived`` fires once ``expected``
    distinct scope warm-checks are simultaneously suspended here, which can only
    happen if the per-scope prepares are gathered rather than awaited one after
    another.
    """

    def __init__(self, expected_concurrency: int):
        self.expected = expected_concurrency
        self.in_flight = 0
        self.max_in_flight = 0
        self.all_arrived = asyncio.Event()
        self.release = asyncio.Event()
        self.keys: list[str] = []

    async def async_get_cache(self, *, key, **kwargs):
        self.in_flight += 1
        self.max_in_flight = max(self.max_in_flight, self.in_flight)
        self.keys.append(key)
        if self.in_flight >= self.expected:
            self.all_arrived.set()
        if not self.release.is_set():
            await self.release.wait()
        self.in_flight -= 1
        return 1.0


@pytest.mark.asyncio
async def test_increment_spend_counters_runs_scopes_concurrently(monkeypatch):
    """The six independent scopes (key, team, team_member, user, end_user+tags,
    org) must prepare their increments concurrently. The probe only fires once
    all eight warm-check reads (one per counter: 6 scopes + 2 tags) are
    suspended in async_get_cache at the same time, which is impossible if the
    awaits are chained sequentially."""
    probe = _ConcurrencyProbe(expected_concurrency=8)
    fake_cache = _make_spend_counter_cache()
    fake_cache.redis_cache.async_get_cache = probe.async_get_cache
    recorded: dict[str, float] = {}

    async def _record_pipeline(increment_list, **_):
        results = []
        for op in increment_list:
            recorded[op["key"]] = recorded.get(op["key"], 0.0) + op["increment_value"]
            results.append(recorded[op["key"]])
        return results

    fake_cache.redis_cache.async_increment_pipeline = AsyncMock(side_effect=_record_pipeline)
    fake_user_cache = _make_user_api_key_cache(get_value=None)
    monkeypatch.setattr(ps, "spend_counter_cache", fake_cache)
    monkeypatch.setattr(ps, "user_api_key_cache", fake_user_cache)
    monkeypatch.setattr(ps, "prisma_client", None)
    monkeypatch.setattr(ps.SpendCounterReseed, "coalesced", AsyncMock(return_value=None))

    task = asyncio.create_task(
        ps.increment_spend_counters(
            token="hashed-tok",
            team_id="t1",
            user_id="u1",
            org_id="org1",
            end_user_id="eu1",
            tags=["a", "b"],
            response_cost=5.0,
        )
    )

    try:
        await asyncio.wait_for(probe.all_arrived.wait(), timeout=2.0)
    except asyncio.TimeoutError:
        probe.release.set()
        await task
        pytest.fail(
            "scope prepares did not run concurrently; sequential awaits "
            f"detected (peak in-flight was {probe.max_in_flight}, expected 8)"
        )

    assert probe.in_flight == 8
    assert probe.max_in_flight == 8
    probe.release.set()
    await task

    assert recorded == {
        "spend:key:hashed-tok": 5.0,
        "spend:team:t1": 5.0,
        "spend:team_member:u1:t1": 5.0,
        "spend:user:u1": 5.0,
        "spend:end_user:eu1": 5.0,
        "spend:tag:a": 5.0,
        "spend:tag:b": 5.0,
        "spend:org:org1": 5.0,
    }


@pytest.mark.asyncio
async def test_increment_spend_counters_skips_reserved_counter_keys(monkeypatch):
    """Counters already reserved by a budget reservation are skipped, every
    other scope is still incremented exactly once, and the reservation is
    finalized after the gathered work completes."""
    import litellm.proxy.spend_tracking.budget_reservation as br

    reserved = {"spend:key:hashed-tok", "spend:org:org1"}
    monkeypatch.setattr(br, "get_reserved_counter_keys", MagicMock(return_value=set(reserved)))
    monkeypatch.setattr(br, "reconcile_budget_reservation", AsyncMock(return_value=()))

    recorded: dict[str, float] = {}

    async def _record_pipeline(increment_list, **_):
        results = []
        for op in increment_list:
            recorded[op["key"]] = recorded.get(op["key"], 0.0) + op["increment_value"]
            results.append(recorded[op["key"]])
        return results

    fake_cache = _make_spend_counter_cache(redis_get_value=None)
    fake_cache.redis_cache.async_increment_pipeline = AsyncMock(side_effect=_record_pipeline)
    fake_user_cache = _make_user_api_key_cache(get_value=None)
    monkeypatch.setattr(ps, "spend_counter_cache", fake_cache)
    monkeypatch.setattr(ps, "user_api_key_cache", fake_user_cache)
    monkeypatch.setattr(ps, "prisma_client", None)
    monkeypatch.setattr(ps.SpendCounterReseed, "coalesced", AsyncMock(return_value=None))

    reservation = {"finalized": False}
    await ps.increment_spend_counters(
        token="hashed-tok",
        team_id="t1",
        user_id="u1",
        org_id="org1",
        end_user_id="eu1",
        tags=["a"],
        response_cost=5.0,
        budget_reservation=reservation,
    )

    assert reservation["finalized"] is True
    assert recorded == {
        "spend:team:t1": 5.0,
        "spend:team_member:u1:t1": 5.0,
        "spend:user:u1": 5.0,
        "spend:end_user:eu1": 5.0,
        "spend:tag:a": 5.0,
    }


@pytest.mark.asyncio
async def test_increment_spend_counters_failing_scope_propagates_after_siblings_settle(
    monkeypatch,
):
    """A failure in one scope must propagate to the caller (so it can invalidate
    reserved counters) while every other scope still settles rather than being
    left as an orphaned background task, and the reservation is not finalized.
    The surviving scopes' increments are still applied in the single pipeline:
    dropping them would under-count spend, the unsafe direction for budget
    enforcement."""
    warmed_keys: list[str] = []

    async def _warm_check(*, key, **kwargs):
        warmed_keys.append(key)
        if key == "spend:team:t1":
            raise RuntimeError("redis get failed")
        return 1.0

    async def _reseed_fails(*, counter_key, **kwargs):
        if counter_key == "spend:team:t1":
            raise RuntimeError("reseed failed")

    applied: dict[str, float] = {}

    async def _record_pipeline(increment_list, **_):
        results = []
        for op in increment_list:
            applied[op["key"]] = op["increment_value"]
            results.append(op["increment_value"])
        return results

    fake_cache = _make_spend_counter_cache()
    fake_cache.redis_cache.async_get_cache = AsyncMock(side_effect=_warm_check)
    fake_cache.redis_cache.async_increment_pipeline = AsyncMock(side_effect=_record_pipeline)
    fake_user_cache = _make_user_api_key_cache(get_value=None)
    monkeypatch.setattr(ps, "spend_counter_cache", fake_cache)
    monkeypatch.setattr(ps, "user_api_key_cache", fake_user_cache)
    monkeypatch.setattr(ps, "prisma_client", None)
    monkeypatch.setattr(
        ps.SpendCounterReseed,
        "coalesced",
        AsyncMock(side_effect=_reseed_fails),
    )

    reservation = {"finalized": False}
    with pytest.raises(RuntimeError, match="reseed failed"):
        await ps.increment_spend_counters(
            token="hashed-tok",
            team_id="t1",
            user_id="u1",
            org_id="org1",
            end_user_id="eu1",
            tags=["a"],
            response_cost=5.0,
            budget_reservation=reservation,
        )

    assert reservation["finalized"] is False
    # every sibling scope settled (its warm-check ran) before the error propagated
    assert set(warmed_keys) == {
        "spend:key:hashed-tok",
        "spend:team:t1",
        "spend:team_member:u1:t1",
        "spend:user:u1",
        "spend:end_user:eu1",
        "spend:tag:a",
        "spend:org:org1",
    }
    # the surviving scopes' increments were still applied, in one pipeline call
    fake_cache.redis_cache.async_increment_pipeline.assert_awaited_once()
    assert applied == {
        "spend:key:hashed-tok": 5.0,
        "spend:team_member:u1:t1": 5.0,
        "spend:user:u1": 5.0,
        "spend:end_user:eu1": 5.0,
        "spend:tag:a": 5.0,
        "spend:org:org1": 5.0,
    }
    fake_cache.redis_cache.async_increment.assert_not_awaited()


@pytest.mark.asyncio
async def test_increment_spend_counters_zero_cost_is_noop_finalizes_reservation(
    monkeypatch,
):
    fake_cache = _make_spend_counter_cache()
    fake_user_cache = _make_user_api_key_cache()
    monkeypatch.setattr(ps, "spend_counter_cache", fake_cache)
    monkeypatch.setattr(ps, "user_api_key_cache", fake_user_cache)
    monkeypatch.setattr(ps, "prisma_client", None)
    reservation = {"finalized": False}

    await ps.increment_spend_counters(
        token="hashed-tok",
        team_id="t1",
        user_id="u1",
        response_cost=0,
        budget_reservation=reservation,
    )

    assert reservation == {"finalized": True}
    assert fake_cache.redis_cache.async_increment.called is False
    fake_cache.redis_cache.async_increment_pipeline.assert_not_awaited()


@pytest.mark.asyncio
async def test_increment_spend_counters_pipelines_all_scopes_in_one_redis_call(
    monkeypatch,
):
    """Every scope's increment must go out in a single async_increment_pipeline
    call, not one INCRBYFLOAT round-trip per scope."""
    counter_cache = ps.DualCache()
    fake_redis = AsyncMock()
    fake_redis.async_get_cache = AsyncMock(return_value=1.0)  # counters warm

    async def _pipeline(increment_list, **_):
        return [1.5] * len(increment_list)

    fake_redis.async_increment_pipeline = AsyncMock(side_effect=_pipeline)
    fake_redis.async_increment = AsyncMock()
    fake_redis.get_ttl = MagicMock(return_value=None)
    counter_cache.redis_cache = fake_redis
    monkeypatch.setattr(ps, "spend_counter_cache", counter_cache)
    monkeypatch.setattr(ps, "user_api_key_cache", ps.DualCache())
    monkeypatch.setattr(ps, "prisma_client", None)

    await ps.increment_spend_counters(
        token="hashed",
        team_id="team-1",
        user_id="user-1",
        response_cost=0.5,
        org_id="org-1",
        end_user_id="eu-1",
        tags=["tag-a", "tag-b"],
    )

    fake_redis.async_increment_pipeline.assert_awaited_once()
    assert fake_redis.async_increment.await_count == 0
    increment_list = fake_redis.async_increment_pipeline.await_args.kwargs["increment_list"]
    expected_keys = {
        "spend:key:hashed",
        "spend:team:team-1",
        "spend:team_member:user-1:team-1",
        "spend:user:user-1",
        "spend:end_user:eu-1",
        "spend:tag:tag-a",
        "spend:tag:tag-b",
        "spend:org:org-1",
    }
    assert {op["key"] for op in increment_list} == expected_keys
    assert all(op["increment_value"] == 0.5 for op in increment_list)
    for key in expected_keys:
        assert counter_cache.in_memory_cache.get_cache(key=key) == 1.5


_PENDING_SPEND_COUNTER_KEYS: tuple[str, ...] = (
    "spend:key:hashed",
    "spend:team:team-1",
    "spend:team_member:user-1:team-1",
    "spend:user:user-1",
    "spend:end_user:eu-1",
    "spend:tag:tag-a",
    "spend:tag:tag-b",
    "spend:org:org-1",
)


def _warm_counters_behind_a_failing_pipeline(monkeypatch, pipeline_error: Exception) -> AsyncMock:
    counter_cache = ps.DualCache()
    for key in _PENDING_SPEND_COUNTER_KEYS:
        counter_cache.in_memory_cache.set_cache(key=key, value=1.0)
    fake_redis = AsyncMock()
    fake_redis.async_get_cache = AsyncMock(return_value=1.0)  # counters warm
    fake_redis.async_increment_pipeline = AsyncMock(side_effect=pipeline_error)
    fake_redis.async_increment = AsyncMock()
    fake_redis.async_delete_cache = AsyncMock()
    fake_redis.get_ttl = MagicMock(return_value=None)
    counter_cache.redis_cache = fake_redis
    monkeypatch.setattr(ps, "spend_counter_cache", counter_cache)
    monkeypatch.setattr(ps, "user_api_key_cache", ps.DualCache())
    monkeypatch.setattr(ps, "prisma_client", None)
    return fake_redis


async def _increment_every_spend_counter() -> None:
    await ps.increment_spend_counters(
        token="hashed",
        team_id="team-1",
        user_id="user-1",
        response_cost=0.5,
        org_id="org-1",
        end_user_id="eu-1",
        tags=["tag-a", "tag-b"],
    )


def _assert_every_pending_counter_invalidated(fake_redis: AsyncMock) -> None:
    assert fake_redis.async_increment.await_count == 0
    deleted_keys = {call.kwargs["key"] for call in fake_redis.async_delete_cache.await_args_list}
    assert deleted_keys == set(_PENDING_SPEND_COUNTER_KEYS)
    for key in _PENDING_SPEND_COUNTER_KEYS:
        assert ps.spend_counter_cache.in_memory_cache.get_cache(key=key) is None


@pytest.mark.asyncio
async def test_increment_spend_counters_pipeline_failure_invalidates_all_counters(
    monkeypatch,
):
    """A failing pipeline must invalidate every pending counter so the next
    request reseeds from the DB (which already holds this request's cost)
    instead of trusting a value the write may have partially applied."""
    from redis.exceptions import ConnectionError as RedisConnectionError

    fake_redis = _warm_counters_behind_a_failing_pipeline(monkeypatch, RedisConnectionError("Redis is down"))

    with pytest.raises(RedisConnectionError):
        await _increment_every_spend_counter()

    _assert_every_pending_counter_invalidated(fake_redis)


@pytest.mark.asyncio
async def test_increment_spend_counters_pool_exhaustion_invalidates_all_counters_without_failing(
    monkeypatch,
):
    """A cluster node with every connection busy raises MaxConnectionsError at once. That is a busy pool, the
    same as the standalone pool's wait timeout, so the counters are invalidated and the request still succeeds."""
    from redis import exceptions as redis_exceptions

    max_connections_error = cast(  # cast-ok: types-redis 4.6 stubs predate MaxConnectionsError
        "type[Exception]",
        redis_exceptions.MaxConnectionsError,  # pyright: ignore[reportAttributeAccessIssue]  # old stubs
    )
    fake_redis = _warm_counters_behind_a_failing_pipeline(monkeypatch, max_connections_error())

    await _increment_every_spend_counter()

    _assert_every_pending_counter_invalidated(fake_redis)


# ---------------------------------------------------------------------------
# _reconcile_budget_reservation_for_counter_update
# ---------------------------------------------------------------------------


@pytest.mark.asyncio
async def test_reconcile_budget_reservation_for_counter_update_returns_empty_set_when_none():
    result = await ps._reconcile_budget_reservation_for_counter_update(budget_reservation=None, response_cost=1.0)
    assert result.reserved_counter_keys == frozenset()
    assert result.pending == ()


@pytest.mark.asyncio
async def test_reconcile_budget_reservation_for_counter_update_failure_invalidates(
    monkeypatch,
):
    """Reservation reconcile raising must invalidate reserved counters, swallow
    the exception, and return an empty set so the caller falls back to the
    direct spend-counter increment instead of skipping it."""
    import litellm.proxy.spend_tracking.budget_reservation as br

    monkeypatch.setattr(
        br,
        "get_reserved_counter_keys",
        MagicMock(return_value={"spend:key:abc"}),
    )
    monkeypatch.setattr(
        br,
        "reconcile_budget_reservation",
        AsyncMock(side_effect=RuntimeError("boom")),
    )
    fake_invalidate = AsyncMock()
    monkeypatch.setattr(br, "invalidate_budget_reservation_counters", fake_invalidate)

    result = await ps._reconcile_budget_reservation_for_counter_update(
        budget_reservation={"foo": "bar"}, response_cost=1.0
    )

    assert result.reserved_counter_keys == frozenset()
    assert result.pending == ()
    assert fake_invalidate.called is True


@pytest.mark.asyncio
async def test_reconcile_budget_reservation_for_counter_update_finalized_reservation_falls_back_to_direct_increment(
    monkeypatch,
):
    """A reservation already finalized before the counter update (the pre-persist
    reconcile failed and dropped its counters) must not shield its keys from the
    direct increment, or the settled cost is never added back after the drop."""
    import litellm.proxy.spend_tracking.budget_reservation as br

    fake_reconcile = AsyncMock()
    monkeypatch.setattr(br, "reconcile_budget_reservation", fake_reconcile)

    result = await ps._reconcile_budget_reservation_for_counter_update(
        budget_reservation={
            "finalized": True,
            "entries": [{"counter_key": "spend:key:abc"}],
        },
        response_cost=1.0,
    )

    assert result.reserved_counter_keys == frozenset()
    assert result.pending == ()
    fake_reconcile.assert_not_awaited()


# ---------------------------------------------------------------------------
# _prepare_end_user_and_tag_spend_increments
# ---------------------------------------------------------------------------


@pytest.mark.asyncio
async def test_prepare_end_user_and_tag_spend_increments_returns_each_unique_tag(
    monkeypatch,
):
    fake_cache = _make_spend_counter_cache(redis_get_value=1.0)
    fake_user_cache = _make_user_api_key_cache()
    monkeypatch.setattr(ps, "spend_counter_cache", fake_cache)
    monkeypatch.setattr(ps, "user_api_key_cache", fake_user_cache)
    monkeypatch.setattr(ps, "prisma_client", None)
    monkeypatch.setattr(ps.SpendCounterReseed, "coalesced", AsyncMock(return_value=None))

    pending = await ps._prepare_end_user_and_tag_spend_increments(
        end_user_id="eu1",
        tags=["a", "b", "a", "", None],
        response_cost=3.0,
        reserved_counter_keys=set(),
    )

    assert {item.counter_key for item in pending} == {
        "spend:end_user:eu1",
        "spend:tag:a",
        "spend:tag:b",
    }
    assert all(item.increment == 3.0 for item in pending)


@pytest.mark.asyncio
async def test_prepare_end_user_and_tag_spend_increments_no_end_user_no_tags_invalid_input_noop(
    monkeypatch,
):
    fake_cache = _make_spend_counter_cache()
    monkeypatch.setattr(ps, "spend_counter_cache", fake_cache)

    pending = await ps._prepare_end_user_and_tag_spend_increments(
        end_user_id=None,
        tags=None,
        response_cost=1.0,
        reserved_counter_keys=set(),
    )

    assert pending == ()
    assert fake_cache.redis_cache.async_increment.called is False


# ---------------------------------------------------------------------------
# _prepare_org_spend_increment
# ---------------------------------------------------------------------------


@pytest.mark.asyncio
async def test_prepare_org_spend_increment_returns_pending_when_org_present(monkeypatch):
    fake_cache = _make_spend_counter_cache(redis_get_value=1.0)
    fake_user_cache = _make_user_api_key_cache()
    monkeypatch.setattr(ps, "spend_counter_cache", fake_cache)
    monkeypatch.setattr(ps, "user_api_key_cache", fake_user_cache)
    monkeypatch.setattr(ps, "prisma_client", None)
    monkeypatch.setattr(ps.SpendCounterReseed, "coalesced", AsyncMock(return_value=None))

    pending = await ps._prepare_org_spend_increment(
        org_id="org-1",
        response_cost=10.0,
        reserved_counter_keys=set(),
    )

    assert len(pending) == 1
    assert pending[0].counter_key == "spend:org:org-1"
    assert pending[0].increment == 10.0


@pytest.mark.asyncio
async def test_prepare_org_spend_increment_no_org_is_noop_invalid_id(monkeypatch):
    fake_cache = _make_spend_counter_cache()
    monkeypatch.setattr(ps, "spend_counter_cache", fake_cache)

    pending = await ps._prepare_org_spend_increment(
        org_id=None,
        response_cost=1.0,
        reserved_counter_keys=set(),
    )

    assert pending == ()
    assert fake_cache.redis_cache.async_increment.called is False


# ---------------------------------------------------------------------------
# _prepare_unreserved_spend_counter_increment
# ---------------------------------------------------------------------------


@pytest.mark.asyncio
async def test_prepare_unreserved_spend_counter_increment_skips_reserved_keys(
    monkeypatch,
):
    fake_cache = _make_spend_counter_cache()
    monkeypatch.setattr(ps, "spend_counter_cache", fake_cache)

    pending = await ps._prepare_unreserved_spend_counter_increment(
        counter_key="spend:tag:x",
        source_cache_key="tag:x",
        increment=1.0,
        reserved_counter_keys={"spend:tag:x"},
    )

    assert pending is None
    assert fake_cache.redis_cache.async_increment.called is False


@pytest.mark.asyncio
async def test_prepare_unreserved_spend_counter_increment_proceeds_when_not_reserved(
    monkeypatch,
):
    fake_cache = _make_spend_counter_cache(redis_get_value=None)
    fake_user_cache = _make_user_api_key_cache()
    reseed = AsyncMock(return_value=None)
    monkeypatch.setattr(ps, "spend_counter_cache", fake_cache)
    monkeypatch.setattr(ps, "user_api_key_cache", fake_user_cache)
    monkeypatch.setattr(ps, "prisma_client", None)
    monkeypatch.setattr(ps.SpendCounterReseed, "coalesced", reseed)

    pending = await ps._prepare_unreserved_spend_counter_increment(
        counter_key="spend:tag:y",
        source_cache_key="tag:y",
        increment=2.0,
        reserved_counter_keys=set(),
    )

    assert pending is not None
    assert pending.counter_key == "spend:tag:y"
    assert pending.increment == 2.0
    assert fake_cache.redis_cache.async_get_cache.called is True
    assert reseed.called is True


# ---------------------------------------------------------------------------
# _prepare_spend_counter_increment
# ---------------------------------------------------------------------------


@pytest.mark.asyncio
async def test_prepare_spend_counter_increment_warm_cache_skips_reseed(monkeypatch):
    fake_cache = _make_spend_counter_cache(redis_get_value=11.0)
    fake_user_cache = _make_user_api_key_cache()
    monkeypatch.setattr(ps, "spend_counter_cache", fake_cache)
    monkeypatch.setattr(ps, "user_api_key_cache", fake_user_cache)
    monkeypatch.setattr(ps, "prisma_client", None)
    reseed = AsyncMock(return_value=None)
    monkeypatch.setattr(ps.SpendCounterReseed, "coalesced", reseed)

    pending = await ps._prepare_spend_counter_increment(
        counter_key="spend:key:k",
        source_cache_key="k",
        increment=3.0,
    )

    assert pending.counter_key == "spend:key:k"
    assert pending.increment == 3.0
    observed = {
        "reseed_called": reseed.called,
        "increment_called": fake_cache.redis_cache.async_increment.called,
        "in_memory_seeded_from_redis": fake_cache.in_memory_cache.set_cache.called,
    }
    assert normalize(observed) == {
        "reseed_called": False,
        "increment_called": False,
        "in_memory_seeded_from_redis": True,
    }


# ---------------------------------------------------------------------------
# _prepare_window_spend_counter_increment
# ---------------------------------------------------------------------------


@pytest.mark.asyncio
async def test_prepare_window_spend_counter_increment_returns_pending_when_initialized(
    monkeypatch,
):
    fake_cache = _make_spend_counter_cache(redis_get_value=0.0)
    monkeypatch.setattr(ps, "spend_counter_cache", fake_cache)
    monkeypatch.setattr(ps, "prisma_client", None)
    monkeypatch.setattr(
        ps.SpendCounterReseed,
        "coalesced_window",
        AsyncMock(return_value=0.0),
    )

    pending = await ps._prepare_window_spend_counter_increment(
        counter_key="spend:key:k:window:1d",
        entity_type="Key",
        entity_id="k",
        window_duration="1d",
        window_start=datetime(2024, 1, 1),
        increment=5.0,
    )

    assert pending is not None
    assert pending.counter_key == "spend:key:k:window:1d"
    assert pending.increment == 5.0


@pytest.mark.asyncio
async def test_prepare_window_spend_counter_increment_missing_window_start_invalid_skips(
    monkeypatch,
):
    fake_cache = _make_spend_counter_cache()
    monkeypatch.setattr(ps, "spend_counter_cache", fake_cache)

    pending = await ps._prepare_window_spend_counter_increment(
        counter_key="spend:key:k:window:1d",
        entity_type="Key",
        entity_id="k",
        window_duration="1d",
        window_start=None,
        increment=5.0,
    )

    assert pending is None
    assert fake_cache.redis_cache.async_increment.called is False


# ---------------------------------------------------------------------------
# _apply_spend_counter_increments
# ---------------------------------------------------------------------------


def _two_pending_increments() -> tuple[ps.PendingSpendIncrement, ...]:
    return (
        ps.PendingSpendIncrement(counter_key="spend:key:k", increment=1.5),
        ps.PendingSpendIncrement(counter_key="spend:team:t", increment=1.5),
    )


@pytest.mark.asyncio
async def test_apply_spend_counter_increments_open_breaker_invalidates_and_returns(monkeypatch):
    """An open Redis circuit breaker is a known, already-logged state, not a per-request tracking failure.

    Re-raising the refusal sent every request through the cost callback's error path, which
    logged an ERROR and fired the failed-tracking alert once per request for the whole outage.
    """
    from litellm.caching.redis_cache import RedisCircuitBreakerOpenError

    fake_cache = _make_spend_counter_cache()
    fake_cache.redis_cache.async_increment_pipeline = AsyncMock(
        side_effect=RedisCircuitBreakerOpenError("Redis circuit breaker is open")
    )
    monkeypatch.setattr(ps, "spend_counter_cache", fake_cache)

    await ps._apply_spend_counter_increments(_two_pending_increments())

    deleted_keys = sorted(call.kwargs["key"] for call in fake_cache.in_memory_cache.delete_cache.call_args_list)
    assert deleted_keys == ["spend:key:k", "spend:team:t"]
    fake_cache.in_memory_cache.set_cache.assert_not_called()


@pytest.mark.asyncio
async def test_apply_spend_counter_increments_redis_timeout_invalidates_and_returns(monkeypatch):
    """A Redis timeout invalidates the counters and returns without reaching the cost callback's error path."""
    from redis.exceptions import TimeoutError as RedisTimeoutError

    fake_cache = _make_spend_counter_cache()
    fake_cache.redis_cache.async_increment_pipeline = AsyncMock(
        side_effect=RedisTimeoutError("Timeout reading from 127.0.0.1:6379")
    )
    monkeypatch.setattr(ps, "spend_counter_cache", fake_cache)

    await ps._apply_spend_counter_increments(_two_pending_increments())

    deleted_keys = sorted(call.kwargs["key"] for call in fake_cache.in_memory_cache.delete_cache.call_args_list)
    assert deleted_keys == ["spend:key:k", "spend:team:t"]
    fake_cache.in_memory_cache.set_cache.assert_not_called()


@pytest.mark.asyncio
async def test_apply_spend_counter_increments_other_redis_error_invalidates_and_raises(monkeypatch):
    fake_cache = _make_spend_counter_cache()
    fake_cache.redis_cache.async_increment_pipeline = AsyncMock(side_effect=ConnectionError("redis down"))
    monkeypatch.setattr(ps, "spend_counter_cache", fake_cache)

    with pytest.raises(ConnectionError, match="redis down"):
        await ps._apply_spend_counter_increments(_two_pending_increments())

    deleted_keys = sorted(call.kwargs["key"] for call in fake_cache.in_memory_cache.delete_cache.call_args_list)
    assert deleted_keys == ["spend:key:k", "spend:team:t"]
    fake_cache.in_memory_cache.set_cache.assert_not_called()


# ---------------------------------------------------------------------------
# _ensure_spend_counter_initialized
# ---------------------------------------------------------------------------


@pytest.mark.asyncio
async def test_ensure_spend_counter_initialized_warm_skips_reseed_and_source(
    monkeypatch,
):
    fake_cache = _make_spend_counter_cache(redis_get_value=20.0)
    fake_user_cache = _make_user_api_key_cache()
    monkeypatch.setattr(ps, "spend_counter_cache", fake_cache)
    monkeypatch.setattr(ps, "user_api_key_cache", fake_user_cache)
    monkeypatch.setattr(ps, "prisma_client", None)
    reseed = AsyncMock(return_value=None)
    monkeypatch.setattr(ps.SpendCounterReseed, "coalesced", reseed)

    await ps._ensure_spend_counter_initialized(
        counter_key="spend:user:u",
        source_cache_key="u",
    )

    observed = {
        "warm_check_redis": fake_cache.redis_cache.async_get_cache.called,
        "reseed_called": reseed.called,
        "source_cache_called": fake_user_cache.async_get_cache.called,
    }
    assert normalize(observed) == {
        "warm_check_redis": True,
        "reseed_called": False,
        "source_cache_called": False,
    }


@pytest.mark.asyncio
async def test_ensure_spend_counter_initialized_cold_seeds_from_source_cache(
    monkeypatch,
):
    fake_cache = _make_spend_counter_cache(redis_get_value=None, redis_increment_value=7.0)
    fake_user_cache = _make_user_api_key_cache(get_value={"spend": 7.0})
    monkeypatch.setattr(ps, "spend_counter_cache", fake_cache)
    monkeypatch.setattr(ps, "user_api_key_cache", fake_user_cache)
    monkeypatch.setattr(ps, "prisma_client", None)
    monkeypatch.setattr(ps.SpendCounterReseed, "coalesced", AsyncMock(return_value=None))

    await ps._ensure_spend_counter_initialized(
        counter_key="spend:user:u",
        source_cache_key="u",
    )

    observed = {
        "source_cache_called": fake_user_cache.async_get_cache.called,
        "seed_increment_called": fake_cache.redis_cache.async_increment.called,
        "warm_check_done": fake_cache.redis_cache.async_get_cache.called,
    }
    assert normalize(observed) == {
        "source_cache_called": True,
        "seed_increment_called": True,
        "warm_check_done": True,
    }


# ---------------------------------------------------------------------------
# _get_source_cache_base_spend
# ---------------------------------------------------------------------------


@pytest.mark.asyncio
async def test_get_source_cache_base_spend_reads_first_hit_from_list(monkeypatch):
    fake_user_cache = MagicMock()

    async def _get(key, **kwargs):
        if key == "miss":
            return None
        if key == "hit-obj":
            obj = MagicMock()
            obj.spend = 12.0
            return obj
        return None

    fake_user_cache.async_get_cache = AsyncMock(side_effect=_get)
    monkeypatch.setattr(ps, "user_api_key_cache", fake_user_cache)

    result = await ps._get_source_cache_base_spend(source_cache_key=["miss", "hit-obj", "miss2"])

    observed = {
        "result": result,
        "calls": fake_user_cache.async_get_cache.call_count,
        "stopped_after_hit": fake_user_cache.async_get_cache.call_count == 2,
    }
    assert normalize(observed) == {
        "result": 12.0,
        "calls": 2,
        "stopped_after_hit": True,
    }


@pytest.mark.asyncio
async def test_get_source_cache_base_spend_no_hits_returns_zero_fallback(monkeypatch):
    """All cache lookups miss — function falls back to 0.0 (no error)."""
    fake_user_cache = _make_user_api_key_cache(get_value=None)
    monkeypatch.setattr(ps, "user_api_key_cache", fake_user_cache)

    result = await ps._get_source_cache_base_spend(source_cache_key="missing-key")
    assert result == 0.0


# ---------------------------------------------------------------------------
# _ensure_window_spend_counter_initialized
# ---------------------------------------------------------------------------


@pytest.mark.asyncio
async def test_ensure_window_spend_counter_initialized_warm_returns_true(monkeypatch):
    fake_cache = _make_spend_counter_cache(redis_get_value=3.0)
    monkeypatch.setattr(ps, "spend_counter_cache", fake_cache)
    monkeypatch.setattr(ps, "prisma_client", None)
    window_reseed = AsyncMock(return_value=0.0)
    monkeypatch.setattr(ps.SpendCounterReseed, "coalesced_window", window_reseed)

    initialized = await ps._ensure_window_spend_counter_initialized(
        counter_key="spend:key:k:window:1d",
        entity_type="Key",
        entity_id="k",
        window_duration="1d",
        window_start=datetime(2024, 1, 1),
    )

    observed = {
        "initialized": initialized,
        "reseed_called": window_reseed.called,
        "redis_get_called": fake_cache.redis_cache.async_get_cache.called,
    }
    assert normalize(observed) == {
        "initialized": True,
        "reseed_called": False,
        "redis_get_called": True,
    }


@pytest.mark.asyncio
async def test_ensure_window_spend_counter_initialized_db_failure_invalid_returns_false(
    monkeypatch,
):
    fake_cache = _make_spend_counter_cache(redis_get_value=None)
    monkeypatch.setattr(ps, "spend_counter_cache", fake_cache)
    monkeypatch.setattr(ps, "prisma_client", None)
    monkeypatch.setattr(
        ps.SpendCounterReseed,
        "coalesced_window",
        AsyncMock(return_value=None),
    )

    initialized = await ps._ensure_window_spend_counter_initialized(
        counter_key="spend:key:k:window:1d",
        entity_type="Key",
        entity_id="k",
        window_duration="1d",
        window_start=datetime(2024, 1, 1),
    )

    assert initialized is False


# ---------------------------------------------------------------------------
# _is_spend_counter_cache_warm
# ---------------------------------------------------------------------------


@pytest.mark.asyncio
async def test_is_spend_counter_cache_warm_redis_hit_seeds_in_memory(monkeypatch):
    fake_cache = _make_spend_counter_cache(redis_get_value=99.0)
    monkeypatch.setattr(ps, "spend_counter_cache", fake_cache)

    result = await ps._is_spend_counter_cache_warm(counter_key="spend:user:u")

    observed = {
        "result": result,
        "redis_get_called": fake_cache.redis_cache.async_get_cache.called,
        "in_memory_set_called": fake_cache.in_memory_cache.set_cache.called,
    }
    assert normalize(observed) == {
        "result": True,
        "redis_get_called": True,
        "in_memory_set_called": True,
    }


@pytest.mark.asyncio
async def test_is_spend_counter_cache_warm_redis_error_falls_back_to_in_memory(
    monkeypatch,
):
    fake_cache = _make_spend_counter_cache(
        redis_get_side_effect=RuntimeError("redis err"),
        in_memory_value=None,
    )
    monkeypatch.setattr(ps, "spend_counter_cache", fake_cache)

    result = await ps._is_spend_counter_cache_warm(counter_key="spend:user:u")
    assert result is False


# ---------------------------------------------------------------------------
# _increment_spend_counter_cache
# ---------------------------------------------------------------------------


@pytest.mark.asyncio
async def test_increment_spend_counter_cache_redis_path_returns_new_value(monkeypatch):
    fake_cache = _make_spend_counter_cache(redis_increment_value=44.0)
    monkeypatch.setattr(ps, "spend_counter_cache", fake_cache)

    result = await ps._increment_spend_counter_cache(counter_key="spend:key:k", increment=4.0)

    observed = {
        "result": result,
        "redis_increment_called": fake_cache.redis_cache.async_increment.called,
        "in_memory_set_called": fake_cache.in_memory_cache.set_cache.called,
    }
    assert normalize(observed) == {
        "result": 44.0,
        "redis_increment_called": True,
        "in_memory_set_called": True,
    }


@pytest.mark.asyncio
async def test_increment_spend_counter_cache_redis_error_raises_and_invalidates(
    monkeypatch,
):
    fake_cache = _make_spend_counter_cache(redis_increment_side_effect=RuntimeError("incr fail"))
    monkeypatch.setattr(ps, "spend_counter_cache", fake_cache)

    with pytest.raises(RuntimeError):
        await ps._increment_spend_counter_cache(counter_key="spend:key:k", increment=1.0)

    assert fake_cache.in_memory_cache.delete_cache.called is True
    assert fake_cache.redis_cache.async_delete_cache.called is True


# ---------------------------------------------------------------------------
# _invalidate_spend_counter
# ---------------------------------------------------------------------------


@pytest.mark.asyncio
async def test_invalidate_spend_counter_deletes_in_memory_and_redis(monkeypatch):
    fake_cache = _make_spend_counter_cache()
    monkeypatch.setattr(ps, "spend_counter_cache", fake_cache)

    await ps._invalidate_spend_counter(counter_key="spend:key:k")

    observed = {
        "in_memory_delete_called": fake_cache.in_memory_cache.delete_cache.called,
        "redis_delete_called": fake_cache.redis_cache.async_delete_cache.called,
        "delete_args_key": fake_cache.redis_cache.async_delete_cache.call_args.kwargs["key"],
    }
    assert normalize(observed) == {
        "in_memory_delete_called": True,
        "redis_delete_called": True,
        "delete_args_key": "spend:key:k",
    }


@pytest.mark.asyncio
async def test_invalidate_spend_counter_swallows_redis_failure_no_raise(monkeypatch):
    fake_cache = _make_spend_counter_cache()
    fake_cache.redis_cache.async_delete_cache = AsyncMock(side_effect=RuntimeError("redis down"))
    monkeypatch.setattr(ps, "spend_counter_cache", fake_cache)

    await ps._invalidate_spend_counter(counter_key="spend:key:k")

    assert fake_cache.in_memory_cache.delete_cache.called is True


# ---------------------------------------------------------------------------
# update_cache
# ---------------------------------------------------------------------------


@pytest.mark.asyncio
async def test_update_cache_user_cache_failure_invalid_state_is_swallowed(monkeypatch):
    """An inner _update_user_cache raising must not propagate — update_cache
    catches and logs, the public coroutine still completes normally."""
    fake_user_cache = MagicMock()
    fake_user_cache.async_batch_get_cache = AsyncMock(side_effect=RuntimeError("cache down"))
    fake_user_cache.async_set_cache_pipeline = AsyncMock()
    monkeypatch.setattr(ps, "user_api_key_cache", fake_user_cache)

    result = await ps.update_cache(
        token=None,
        user_id="u1",
        end_user_id=None,
        team_id=None,
        response_cost=1.0,
        parent_otel_span=None,
        tags=None,
    )

    assert result is None


async def _drain_spend_cache_writes() -> None:
    for _ in range(3):
        await asyncio.sleep(0)


async def _update_team_spend(team_id: str, response_cost: float) -> None:
    await ps.update_cache(
        token=None,
        user_id=None,
        end_user_id=None,
        team_id=team_id,
        response_cost=response_cost,
        parent_otel_span=None,
    )


@pytest.mark.asyncio
async def test_update_cache_adds_spend_to_the_cached_team_only(monkeypatch: pytest.MonkeyPatch) -> None:
    cache: Final = UserApiKeyCache()
    cache.set_cache(key="team_id:warm", value=LiteLLM_TeamTableCachedObj(team_id="warm", spend=5.0, blocked=True))
    monkeypatch.setattr(ps, "user_api_key_cache", cache)

    await _update_team_spend(team_id="warm", response_cost=2.5)
    await _update_team_spend(team_id="cold", response_cost=2.5)
    await _drain_spend_cache_writes()

    warm: Final = cache.get_cache(key="team_id:warm", model_type=LiteLLM_TeamTableCachedObj)
    assert warm is not None
    assert (warm.spend, warm.blocked) == (7.5, True)
    assert cache.get_cache(key="team_id:cold") is None


async def _block_on_this_worker(cache: UserApiKeyCache, team_id: str) -> None:
    await refresh_cached_team(
        team_row=LiteLLM_TeamTable(team_id=team_id, blocked=True),
        user_api_key_cache=cache,
        proxy_logging_obj=ps.proxy_logging_obj,
    )


async def _invalidated_by_a_peer(cache: UserApiKeyCache, team_id: str) -> None:
    evict_local(cache_keys=(f"team_id:{team_id}",), user_api_key_cache=cache)


@pytest.mark.asyncio
@pytest.mark.parametrize(
    ("management_write", "expected_blocked"),
    [(_block_on_this_worker, True), (_invalidated_by_a_peer, None)],
)
async def test_update_cache_never_restores_a_team_a_management_write_replaced(
    management_write: Callable[[UserApiKeyCache, str], Awaitable[None]],
    expected_blocked: bool | None,
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    team_id: Final = "kill-switch"
    cache: Final = UserApiKeyCache()
    cache.set_cache(key=f"team_id:{team_id}", value=LiteLLM_TeamTableCachedObj(team_id=team_id, blocked=False))
    monkeypatch.setattr(ps, "user_api_key_cache", cache)

    await _update_team_spend(team_id=team_id, response_cost=1.0)
    await management_write(cache, team_id)
    await _drain_spend_cache_writes()

    seen_by_auth: Final = cache.get_cache(key=f"team_id:{team_id}", model_type=LiteLLM_TeamTableCachedObj)
    assert (seen_by_auth.blocked if seen_by_auth is not None else None) == expected_blocked


_CachedSpendObject = LiteLLM_UserTable | LiteLLM_EndUserTable | LiteLLM_TeamTableCachedObj | LiteLLM_TagTable


def _cached_spend_object(
    kind: Literal["user", "end_user", "team", "tag"], models: list[str]
) -> tuple[str, _CachedSpendObject]:
    match kind:
        case "user":
            return "race-user", LiteLLM_UserTable(user_id="race-user", models=models)
        case "end_user":
            return end_user_cache_key("race-end-user"), LiteLLM_EndUserTable(
                user_id="race-end-user", blocked=False, models=models
            )
        case "team":
            return "team_id:race-team", LiteLLM_TeamTableCachedObj(team_id="race-team", models=models)
        case "tag":
            return tag_cache_key("race-tag"), LiteLLM_TagTable(tag_name="race-tag", models=models)


class _RunsOnKeyRead(InMemoryCache):
    def __init__(self, on_read: Callable[[], Awaitable[None]]) -> None:
        super().__init__()
        self._on_read: Final = on_read

    async def async_get_cache(self, key: str, **kwargs: object) -> object:
        await self._on_read()
        return await super().async_get_cache(key, **kwargs)


def _cache_that_runs_on_key_read(on_read: Callable[[], Awaitable[None]]) -> UserApiKeyCache:
    return UserApiKeyCache(in_memory_cache=InMemoryCache(), key_object_in_memory_cache=_RunsOnKeyRead(on_read))


@pytest.fixture
def proxy_max_budget(monkeypatch: pytest.MonkeyPatch) -> None:
    monkeypatch.setattr(litellm, "max_budget", 100.0)


@pytest.mark.asyncio
@pytest.mark.parametrize("kind", ["user", "end_user", "team", "tag"])
@pytest.mark.parametrize(
    ("management_write", "expected"),
    [("none", (["gpt-a"], 0.25)), ("evict", None), ("refresh", (["gpt-b"], 0.25))],
)
async def test_update_cache_adds_spend_to_the_entry_a_management_write_left_mid_flight(
    kind: Literal["user", "end_user", "team", "tag"],
    management_write: Literal["none", "evict", "refresh"],
    expected: tuple[list[str], float] | None,
) -> None:
    key, seeded = _cached_spend_object(kind, models=["gpt-a"])
    _, refreshed = _cached_spend_object(kind, models=["gpt-b"])

    async def _write() -> None:
        if management_write == "none":
            return
        evict_local(cache_keys=(key,), user_api_key_cache=cache)
        if management_write == "refresh":
            cache.set_cache(key=key, value=refreshed, local_only=True)

    cache: Final = _cache_that_runs_on_key_read(on_read=_write)
    cache.set_cache(key=key, value=seeded, local_only=True)

    await ps.update_cache(
        token="sk-race",
        user_id="race-user",
        end_user_id="race-end-user",
        team_id="race-team",
        response_cost=0.25,
        parent_otel_span=None,
        tags=["race-tag"],
        cache=cache,
    )
    await _drain_spend_cache_writes()

    stored: Final = cache.get_cache(key=key, local_only=True, model_type=type(seeded))
    assert ((stored.models, stored.spend) if stored is not None else None) == expected


@pytest.mark.asyncio
@pytest.mark.usefixtures("proxy_max_budget")
@pytest.mark.parametrize("user", ["cached_on_this_worker", "not_cached_on_this_worker", "none"])
async def test_update_cache_adds_spend_to_the_global_proxy_spend_for_every_request(
    user: Literal["cached_on_this_worker", "not_cached_on_this_worker", "none"],
) -> None:
    cache: Final = UserApiKeyCache()
    cache.set_cache(key=GLOBAL_PROXY_SPEND_CACHE_KEY, value=10.0)
    if user == "cached_on_this_worker":
        cache.set_cache(key="spender", value=LiteLLM_UserTable(user_id="spender", spend=1.0))

    await ps.update_cache(
        token=None,
        user_id=None if user == "none" else "spender",
        end_user_id=None,
        team_id=None,
        response_cost=0.25,
        parent_otel_span=None,
        cache=cache,
    )
    await _drain_spend_cache_writes()

    assert cache.get_cache(key=GLOBAL_PROXY_SPEND_CACHE_KEY) == 10.25


@pytest.mark.asyncio
@pytest.mark.usefixtures("proxy_max_budget")
async def test_update_cache_never_brings_back_a_global_proxy_spend_a_budget_reset_cleared() -> None:
    async def _budget_reset() -> None:
        await cache.async_delete_cache(key=GLOBAL_PROXY_SPEND_CACHE_KEY)

    cache: Final = _cache_that_runs_on_key_read(on_read=_budget_reset)
    cache.set_cache(key=GLOBAL_PROXY_SPEND_CACHE_KEY, value=10.0)

    await ps.update_cache(
        token="sk-spender",
        user_id="spender",
        end_user_id=None,
        team_id=None,
        response_cost=0.25,
        parent_otel_span=None,
        cache=cache,
    )
    await _drain_spend_cache_writes()

    assert cache.get_cache(key=GLOBAL_PROXY_SPEND_CACHE_KEY) is None


class _FakeRedisRunningIncrementIfPresent(FakeAsyncRedis):
    def __init__(self, server: FakeServer, after_increment: Callable[[], object] = lambda: None) -> None:
        super().__init__(server=server)
        self._after_increment: Final = after_increment

    async def eval(self, script: str, numkeys: int, *keys_and_args: str | bytes | float) -> bytes | None:
        key, increment = keys_and_args
        if not await self.exists(key):
            return None
        total: Final = await self.incrbyfloat(key, increment)
        self._after_increment()
        return str(total).encode()


class _RedisCacheOnClient(RedisCache):
    def __init__(self, client: FakeAsyncRedis) -> None:
        self._injected_client: Final = client
        self.namespace = None
        self._circuit_breaker = RedisCircuitBreaker(failure_threshold=5, recovery_timeout=30)
        self.service_logger_obj = ServiceLogging()
        self.default_ttl = None

    def init_async_client(self) -> FakeAsyncRedis:
        return self._injected_client


class _FakeRedisHoldingOneReply(FakeAsyncRedis):
    def __init__(self, server: FakeServer, held_increment: float, applied: asyncio.Event, released: asyncio.Event) -> None:
        super().__init__(server=server)
        self._held_increment: Final = held_increment
        self._applied: Final = applied
        self._released: Final = released

    async def eval(self, script: str, numkeys: int, *keys_and_args: str | bytes | float) -> bytes | None:
        key, increment = keys_and_args
        total: Final = await self.incrbyfloat(key, increment)
        if float(increment) == self._held_increment:
            self._applied.set()
            await self._released.wait()
        return str(total).encode()


_SharedGlobalSpendOutcome = Literal["peers_added_spend", "reset_elsewhere", "reset_while_redis_answers", "redis_down"]


@pytest.mark.asyncio
@pytest.mark.usefixtures("proxy_max_budget")
@pytest.mark.parametrize(
    ("outcome", "shared_total", "expected"),
    [
        ("peers_added_spend", 11.75, 12.0),
        ("reset_elsewhere", None, 10.25),
        ("reset_while_redis_answers", 9.75, None),
        ("redis_down", 10.0, 10.25),
    ],
)
async def test_update_cache_moves_this_workers_global_proxy_spend_to_the_shared_total(
    outcome: _SharedGlobalSpendOutcome, shared_total: float | None, expected: float | None
) -> None:
    def _reset_this_worker() -> None:
        if outcome == "reset_while_redis_answers":
            cache.in_memory_cache.delete_cache(GLOBAL_PROXY_SPEND_CACHE_KEY)

    server: Final = FakeServer()
    redis: Final = _FakeRedisRunningIncrementIfPresent(server, after_increment=_reset_this_worker)
    if shared_total is not None:
        await redis.set(GLOBAL_PROXY_SPEND_CACHE_KEY, shared_total)
    server.connected = outcome != "redis_down"
    cache: Final = UserApiKeyCache(in_memory_cache=InMemoryCache(), redis_cache=_RedisCacheOnClient(redis))
    cache.set_cache(key=GLOBAL_PROXY_SPEND_CACHE_KEY, value=10.0, local_only=True)

    await ps.update_cache(
        token=None,
        user_id=None,
        end_user_id=None,
        team_id=None,
        response_cost=0.25,
        parent_otel_span=None,
        cache=cache,
    )

    assert cache.get_cache(key=GLOBAL_PROXY_SPEND_CACHE_KEY, local_only=True) == expected


@pytest.mark.asyncio
@pytest.mark.usefixtures("proxy_max_budget")
async def test_update_cache_keeps_the_newest_global_proxy_spend_when_redis_replies_out_of_order() -> None:
    held_increment_applied: Final = asyncio.Event()
    held_reply_released: Final = asyncio.Event()
    server: Final = FakeServer()
    redis: Final = _FakeRedisHoldingOneReply(
        server, held_increment=0.25, applied=held_increment_applied, released=held_reply_released
    )
    await redis.set(GLOBAL_PROXY_SPEND_CACHE_KEY, 10.0)
    cache: Final = UserApiKeyCache(in_memory_cache=InMemoryCache(), redis_cache=_RedisCacheOnClient(redis))
    cache.set_cache(key=GLOBAL_PROXY_SPEND_CACHE_KEY, value=10.0, local_only=True)

    async def _request(response_cost: float) -> None:
        await ps.update_cache(
            token=None,
            user_id=None,
            end_user_id=None,
            team_id=None,
            response_cost=response_cost,
            parent_otel_span=None,
            cache=cache,
        )

    earlier: Final = asyncio.create_task(_request(0.25))
    await held_increment_applied.wait()
    await _request(0.5)
    held_reply_released.set()
    await earlier

    assert cache.get_cache(key=GLOBAL_PROXY_SPEND_CACHE_KEY, local_only=True) == 10.75


@pytest.mark.asyncio
@pytest.mark.usefixtures("proxy_max_budget")
async def test_update_cache_adds_the_global_spend_scalar_in_redis_as_spend_counters() -> None:
    target_during_increment: Final = MagicMock()
    redis: Final = _FakeRedisRunningIncrementIfPresent(
        FakeServer(), after_increment=lambda: target_during_increment(current_service_target())
    )
    await redis.set(GLOBAL_PROXY_SPEND_CACHE_KEY, 10.0)
    cache: Final = UserApiKeyCache(in_memory_cache=InMemoryCache(), redis_cache=_RedisCacheOnClient(redis))

    await ps.update_cache(
        token=None,
        user_id="u1",
        end_user_id=None,
        team_id=None,
        response_cost=2.0,
        parent_otel_span=None,
        cache=cache,
    )

    assert target_during_increment.call_args_list == [call("spend_counters")]
    assert float(await redis.get(GLOBAL_PROXY_SPEND_CACHE_KEY)) == 12.0
    assert current_service_target() is None


@pytest.mark.asyncio
@pytest.mark.usefixtures("proxy_max_budget")
@pytest.mark.parametrize(("redis_up", "expected_status"), [(True, StatusCode.OK), (False, StatusCode.ERROR)])
async def test_update_cache_traces_the_global_spend_redis_call_under_the_request_span(
    monkeypatch: pytest.MonkeyPatch, redis_up: bool, expected_status: StatusCode
) -> None:
    exporter: Final = InMemorySpanExporter()
    provider: Final = TracerProvider()
    provider.add_span_processor(SimpleSpanProcessor(exporter))
    monkeypatch.setattr(
        litellm,
        "service_callback",
        [OpenTelemetry(config=OpenTelemetryConfig(exporter=exporter), tracer_provider=provider)],
    )
    server: Final = FakeServer()
    redis: Final = _FakeRedisRunningIncrementIfPresent(server)
    await redis.set(GLOBAL_PROXY_SPEND_CACHE_KEY, 10.0)
    server.connected = redis_up
    cache: Final = UserApiKeyCache(in_memory_cache=InMemoryCache(), redis_cache=_RedisCacheOnClient(redis))
    request_span: Final = provider.get_tracer("test").start_span("request")

    await ps.update_cache(
        token=None,
        user_id=None,
        end_user_id=None,
        team_id=None,
        response_cost=2.0,
        parent_otel_span=request_span,
        cache=cache,
    )
    await _drain_spend_cache_writes()
    request_span.end()

    redis_spans: Final = tuple(
        (span.attributes.get("call_type"), span.status.status_code, span.parent.span_id)
        for span in exporter.get_finished_spans()
        if span.name == "redis" and span.attributes is not None and span.parent is not None
    )
    assert redis_spans == (
        ("async_increment_if_present", expected_status, request_span.get_span_context().span_id),
    )


class _UnwritableKeyUserApiKeyCache(UserApiKeyCache):
    def __init__(self, unwritable_key: str) -> None:
        super().__init__(in_memory_cache=InMemoryCache())
        self._unwritable_key: Final = unwritable_key

    def set_cache(self, key: str | None, value: object, local_only: bool = False, **kwargs: object):
        if key == self._unwritable_key:
            raise RuntimeError("cache write failed")
        return super().set_cache(key, value, local_only=local_only, **kwargs)


@pytest.mark.asyncio
@pytest.mark.usefixtures("proxy_max_budget")
async def test_update_cache_still_adds_spend_to_the_others_when_one_cached_object_fails() -> None:
    cache: Final = _UnwritableKeyUserApiKeyCache(unwritable_key="broken-user")
    cache.in_memory_cache.set_cache(
        key="broken-user",
        value=CacheCodec.serialize(LiteLLM_UserTable(user_id="broken-user", spend=1.0), model_type=LiteLLM_UserTable),
    )
    cache.set_cache(key=GLOBAL_PROXY_SPEND_CACHE_KEY, value=10.0)
    cache.set_cache(key="team_id:still-billed", value=LiteLLM_TeamTableCachedObj(team_id="still-billed", spend=1.0))

    await ps.update_cache(
        token=None,
        user_id="broken-user",
        end_user_id=None,
        team_id="still-billed",
        response_cost=0.25,
        parent_otel_span=None,
        cache=cache,
    )
    await _drain_spend_cache_writes()

    team: Final = cache.get_cache(key="team_id:still-billed", model_type=LiteLLM_TeamTableCachedObj)
    assert (team.spend if team is not None else None, cache.get_cache(key=GLOBAL_PROXY_SPEND_CACHE_KEY)) == (
        1.25,
        10.25,
    )