mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-03 02:22:24 +00:00
test(proxy): add failing regressions for durable spend-log flush behavior
Pin expected re-queue on flush failure, deadlock/pool-timeout retries, and per-request Redis spend-log buffering so finance-report data loss cannot recur without failing tests.
This commit is contained in:
mubashir1osmani
parent
586e31f1c4
commit
9ab0875fc2
2 changed files with 282 additions and 0 deletions
103
tests/test_litellm/proxy/db/test_spend_log_row_redis_buffer.py
Normal file
103
tests/test_litellm/proxy/db/test_spend_log_row_redis_buffer.py
Normal file
|
|
@ -0,0 +1,103 @@
|
|||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
|
||||
import pytest
|
||||
|
||||
from litellm.proxy.db.db_spend_update_writer import DBSpendUpdateWriter
|
||||
|
||||
EXPECTED_REDIS_SPEND_LOG_BUFFER_KEY = "litellm_spend_log_buffer"
|
||||
|
||||
|
||||
def _redis_spend_log_buffer_key() -> str:
|
||||
from litellm.constants import REDIS_SPEND_LOG_BUFFER_KEY
|
||||
|
||||
return REDIS_SPEND_LOG_BUFFER_KEY
|
||||
|
||||
|
||||
def test_db_spend_update_writer_exposes_spend_log_redis_buffer() -> None:
|
||||
writer = DBSpendUpdateWriter(redis_cache=MagicMock())
|
||||
assert hasattr(writer, "spend_log_redis_buffer")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_insert_spend_log_to_db_buffers_memory_and_redis_when_enabled(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
import litellm.proxy.proxy_server as proxy_server_mod
|
||||
|
||||
monkeypatch.setattr(
|
||||
proxy_server_mod,
|
||||
"general_settings",
|
||||
{"use_redis_transaction_buffer": True},
|
||||
)
|
||||
|
||||
mock_redis_cache = MagicMock()
|
||||
mock_redis_cache.async_rpush = AsyncMock()
|
||||
mock_prisma_client = MagicMock()
|
||||
mock_prisma_client.spend_log_transactions = []
|
||||
mock_prisma_client._spend_log_transactions_lock = asyncio.Lock()
|
||||
|
||||
writer = DBSpendUpdateWriter(redis_cache=mock_redis_cache)
|
||||
payload = {
|
||||
"request_id": "req-buffer-test",
|
||||
"spend": 0.42,
|
||||
"model": "gpt-4o-mini",
|
||||
}
|
||||
|
||||
await writer._insert_spend_log_to_db(
|
||||
payload=payload,
|
||||
prisma_client=mock_prisma_client,
|
||||
)
|
||||
|
||||
assert mock_prisma_client.spend_log_transactions == [payload]
|
||||
mock_redis_cache.async_rpush.assert_awaited_once()
|
||||
assert (
|
||||
mock_redis_cache.async_rpush.await_args.args[0]
|
||||
== _redis_spend_log_buffer_key()
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_update_spend_logs_job_flushes_rows_buffered_only_in_redis(
|
||||
mock_prisma_client: object,
|
||||
make_spend_log_row: object,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
from litellm.proxy.utils import update_spend_logs_job
|
||||
|
||||
import litellm.proxy.guardrails.usage_tracking as guard_mod
|
||||
import litellm.proxy.db.spend_log_tool_index as tool_mod
|
||||
|
||||
monkeypatch.setattr(
|
||||
guard_mod, "process_spend_logs_guardrail_usage", AsyncMock(), raising=False
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
tool_mod, "process_spend_logs_tool_usage", AsyncMock(), raising=False
|
||||
)
|
||||
|
||||
mock_prisma_client.spend_log_transactions = [] # type: ignore[attr-defined]
|
||||
mock_prisma_client.db.litellm_spendlogs.create_many = AsyncMock() # type: ignore[attr-defined]
|
||||
|
||||
redis_row = make_spend_log_row(request_id="redis-only") # type: ignore[operator]
|
||||
redis_buffer = MagicMock()
|
||||
redis_buffer.is_enabled = MagicMock(return_value=True)
|
||||
redis_buffer.get_buffered_row_count = AsyncMock(return_value=1)
|
||||
redis_buffer.pop_buffered_spend_log_rows = AsyncMock(return_value=[redis_row])
|
||||
redis_buffer.requeue_spend_log_rows = AsyncMock()
|
||||
|
||||
proxy_logging = MagicMock()
|
||||
proxy_logging.failure_handler = AsyncMock()
|
||||
proxy_logging.db_spend_update_writer = MagicMock()
|
||||
proxy_logging.db_spend_update_writer.spend_log_redis_buffer = redis_buffer
|
||||
|
||||
await update_spend_logs_job(
|
||||
prisma_client=mock_prisma_client, # type: ignore[arg-type]
|
||||
db_writer_client=None,
|
||||
proxy_logging_obj=proxy_logging,
|
||||
)
|
||||
|
||||
assert mock_prisma_client.db.litellm_spendlogs.create_many.await_count == 1 # type: ignore[attr-defined]
|
||||
flushed = mock_prisma_client.db.litellm_spendlogs.create_many.await_args.kwargs["data"][0] # type: ignore[attr-defined]
|
||||
assert flushed["request_id"] == "redis-only"
|
||||
|
|
@ -0,0 +1,179 @@
|
|||
"""Regression tests for durable spend-log flushing.
|
||||
|
||||
These tests pin the behavior finance needs after high-load DB incidents:
|
||||
(1) failed flush batches are re-queued instead of dropped
|
||||
(2) deadlock and pool-timeout errors are retried, not only httpx connection errors
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
|
||||
import httpx
|
||||
import prisma.errors
|
||||
import pytest
|
||||
|
||||
from litellm.proxy.utils import ProxyUpdateSpend, update_spend_logs_job
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"error,expected_retryable",
|
||||
[
|
||||
(httpx.ConnectError("connect"), True),
|
||||
(httpx.ReadError("read"), True),
|
||||
(httpx.PoolTimeout("pool timeout"), True),
|
||||
(prisma.errors.PrismaError("deadlock detected"), True),
|
||||
(prisma.errors.PrismaError("pool timeout waiting for connection"), True),
|
||||
(ValueError("invalid spend log row"), False),
|
||||
],
|
||||
)
|
||||
def test_spend_log_flush_retryable_error_classifier(
|
||||
error: Exception, expected_retryable: bool
|
||||
) -> None:
|
||||
from litellm.proxy._types import is_spend_log_flush_retryable_error
|
||||
|
||||
assert is_spend_log_flush_retryable_error(error) is expected_retryable
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_update_spend_logs_retries_on_deadlock_error(
|
||||
mock_prisma_client: object,
|
||||
make_spend_log_row: object,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
import litellm.proxy.utils as utils_mod
|
||||
|
||||
async def _fake_sleep(_: float) -> None:
|
||||
return None
|
||||
|
||||
monkeypatch.setattr(utils_mod.asyncio, "sleep", _fake_sleep)
|
||||
|
||||
create_many = AsyncMock(
|
||||
side_effect=[
|
||||
prisma.errors.PrismaError("deadlock detected"),
|
||||
None,
|
||||
]
|
||||
)
|
||||
mock_prisma_client.db.litellm_spendlogs.create_many = create_many # type: ignore[attr-defined]
|
||||
proxy_logging = MagicMock()
|
||||
proxy_logging.failure_handler = AsyncMock()
|
||||
logs = [make_spend_log_row(request_id="r1")] # type: ignore[operator]
|
||||
|
||||
await ProxyUpdateSpend.update_spend_logs(
|
||||
n_retry_times=1,
|
||||
prisma_client=mock_prisma_client, # type: ignore[arg-type]
|
||||
db_writer_client=None,
|
||||
proxy_logging_obj=proxy_logging,
|
||||
logs_to_process=logs,
|
||||
)
|
||||
|
||||
assert create_many.await_count == 2
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_update_spend_logs_retries_on_pool_timeout(
|
||||
mock_prisma_client: object,
|
||||
make_spend_log_row: object,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
import litellm.proxy.utils as utils_mod
|
||||
|
||||
async def _fake_sleep(_: float) -> None:
|
||||
return None
|
||||
|
||||
monkeypatch.setattr(utils_mod.asyncio, "sleep", _fake_sleep)
|
||||
|
||||
create_many = AsyncMock(
|
||||
side_effect=[
|
||||
httpx.PoolTimeout("pool timeout"),
|
||||
None,
|
||||
]
|
||||
)
|
||||
mock_prisma_client.db.litellm_spendlogs.create_many = create_many # type: ignore[attr-defined]
|
||||
proxy_logging = MagicMock()
|
||||
proxy_logging.failure_handler = AsyncMock()
|
||||
logs = [make_spend_log_row(request_id="r1")] # type: ignore[operator]
|
||||
|
||||
await ProxyUpdateSpend.update_spend_logs(
|
||||
n_retry_times=1,
|
||||
prisma_client=mock_prisma_client, # type: ignore[arg-type]
|
||||
db_writer_client=None,
|
||||
proxy_logging_obj=proxy_logging,
|
||||
logs_to_process=logs,
|
||||
)
|
||||
|
||||
assert create_many.await_count == 2
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_update_spend_logs_requeues_after_retryable_flush_failure(
|
||||
mock_prisma_client: object,
|
||||
make_spend_log_row: object,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
import litellm.proxy.utils as utils_mod
|
||||
|
||||
async def _fake_sleep(_: float) -> None:
|
||||
return None
|
||||
|
||||
monkeypatch.setattr(utils_mod.asyncio, "sleep", _fake_sleep)
|
||||
|
||||
failed_logs = [make_spend_log_row(request_id="r1")] # type: ignore[operator]
|
||||
mock_prisma_client.db.litellm_spendlogs.create_many = AsyncMock( # type: ignore[attr-defined]
|
||||
side_effect=httpx.ReadError("network blip")
|
||||
)
|
||||
proxy_logging = MagicMock()
|
||||
proxy_logging.failure_handler = AsyncMock()
|
||||
|
||||
with pytest.raises(httpx.ReadError):
|
||||
await ProxyUpdateSpend.update_spend_logs(
|
||||
n_retry_times=1,
|
||||
prisma_client=mock_prisma_client, # type: ignore[arg-type]
|
||||
db_writer_client=None,
|
||||
proxy_logging_obj=proxy_logging,
|
||||
logs_to_process=failed_logs,
|
||||
)
|
||||
|
||||
assert mock_prisma_client.spend_log_transactions == failed_logs # type: ignore[attr-defined]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_update_spend_logs_job_requeues_after_retryable_flush_failure(
|
||||
mock_prisma_client: object,
|
||||
make_spend_log_row: object,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
import litellm.proxy.guardrails.usage_tracking as guard_mod
|
||||
import litellm.proxy.db.spend_log_tool_index as tool_mod
|
||||
import litellm.proxy.utils as utils_mod
|
||||
|
||||
async def _fake_sleep(_: float) -> None:
|
||||
return None
|
||||
|
||||
monkeypatch.setattr(utils_mod.asyncio, "sleep", _fake_sleep)
|
||||
monkeypatch.setattr(
|
||||
guard_mod, "process_spend_logs_guardrail_usage", AsyncMock(), raising=False
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
tool_mod, "process_spend_logs_tool_usage", AsyncMock(), raising=False
|
||||
)
|
||||
|
||||
failed_row = make_spend_log_row(request_id="lost-if-not-requeued") # type: ignore[operator]
|
||||
mock_prisma_client.spend_log_transactions = [failed_row] # type: ignore[attr-defined]
|
||||
mock_prisma_client.db.litellm_spendlogs.create_many = AsyncMock( # type: ignore[attr-defined]
|
||||
side_effect=httpx.ReadError("connection reset")
|
||||
)
|
||||
|
||||
proxy_logging = MagicMock()
|
||||
proxy_logging.failure_handler = AsyncMock()
|
||||
proxy_logging.db_spend_update_writer = MagicMock()
|
||||
proxy_logging.db_spend_update_writer.spend_log_redis_buffer = None
|
||||
|
||||
with pytest.raises(httpx.ReadError):
|
||||
await update_spend_logs_job(
|
||||
prisma_client=mock_prisma_client, # type: ignore[arg-type]
|
||||
db_writer_client=None,
|
||||
proxy_logging_obj=proxy_logging,
|
||||
)
|
||||
|
||||
assert mock_prisma_client.spend_log_transactions == [failed_row] # type: ignore[attr-defined]
|
||||
Loading…
Add table
Reference in a new issue