scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-21 00:21:49 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 74

test(spend_tracking): type the spend log flush helpers and cover the object reads

Browse source 
The fake table, the wiring helpers and the new tests took Any where the fixtures have concrete
types. The collapse test now walks the file, fine-tuning, video, container and Responses reads
and an unknown call type, and a sibling test walks the inference call types that are re-keyed.
This commit is contained in:
Yucheng He 2026-09-15 12:28:13 -07:00
parent ce84bbf81e
commit abd8c08315
1 changed files with 72 additions and 27 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

99
tests/test_litellm/proxy/utils/prisma_and_spend/test_proxy_update_spend.py
View file

@ -13,8 +13,8 @@ import json
import logging
import httpx
from collections.abc import Iterator
from typing import Any, Dict, List
from collections.abc import Callable, Iterator, Mapping, Sequence
from typing import Any, Dict, List, TypeAlias
from unittest.mock import AsyncMock, MagicMock
import pytest
@ -887,28 +887,35 @@ def test_disable_spend_updates_error_when_general_settings_unavailable(
ProxyUpdateSpend.disable_spend_updates()
SpendLogRow: TypeAlias = dict[str, object]
"""One LiteLLM_SpendLogs row as the flush hands it to prisma."""
SpendLogRowFactory: TypeAlias = Callable[..., SpendLogRow]
class _SpendLogsTable:
"""``LiteLLM_SpendLogs`` as the flush sees it: ``create_many`` is ``INSERT ... ON CONFLICT
DO NOTHING`` on ``request_id`` and reports how many rows landed, and ``query_raw`` answers the
identity read-back with the stored ``(request_id, litellm_call_id)`` pairs."""
def __init__(self, seeded: List[Dict[str, Any]] | None = None) -> None:
self.rows: Dict[str, Dict[str, Any]] = {row["request_id"]: dict(row) for row in seeded or []}
self.inserts: List[List[str]] = []
def __init__(self, seeded: Sequence[SpendLogRow] | None = None) -> None:
self.rows: dict[str, SpendLogRow] = {str(row["request_id"]): dict(row) for row in seeded or []}
self.inserts: list[list[str]] = []
self.query_raw_calls = 0
async def create_many(self, *, data: Any, skip_duplicates: bool) -> int:
async def create_many(self, *, data: Sequence[SpendLogRow], skip_duplicates: bool) -> int:
assert skip_duplicates is True
self.inserts.append([row["request_id"] for row in data])
self.inserts.append([str(row["request_id"]) for row in data])
inserted = 0
for row in data:
if row["request_id"] in self.rows:
request_id = str(row["request_id"])
if request_id in self.rows:
continue
self.rows[row["request_id"]] = dict(row)
self.rows[request_id] = dict(row)
inserted += 1
return inserted
async def query_raw(self, sql: str, request_ids: Any) -> List[Dict[str, Any]]:
async def query_raw(self, sql: str, request_ids: Sequence[str]) -> list[SpendLogRow]:
self.query_raw_calls += 1
assert "WHERE request_id = ANY($1::text[])" in sql
return [
@ -918,14 +925,16 @@ class _SpendLogsTable:
]
def _wire_spend_logs_table(mock_prisma_client: Any, seeded: List[Dict[str, Any]] | None = None) -> _SpendLogsTable:
def _wire_spend_logs_table(
mock_prisma_client: MagicMock, seeded: Sequence[SpendLogRow] | None = None
) -> _SpendLogsTable:
table = _SpendLogsTable(seeded)
mock_prisma_client.db.litellm_spendlogs.create_many = AsyncMock(side_effect=table.create_many)
mock_prisma_client.db.query_raw = AsyncMock(side_effect=table.query_raw)
return table
def _wire_routed_db(mock_prisma_client: Any, table: _SpendLogsTable) -> MagicMock:
def _wire_routed_db(mock_prisma_client: MagicMock, table: _SpendLogsTable) -> MagicMock:
"""``prisma_client.db`` as a read-replica router: top-level reads go to the reader, which never
sees this flush's writes, and ``writer`` is the engine the rows landed on."""
routed = MagicMock(spec=RoutingPrismaWrapper)
@ -939,7 +948,7 @@ def _wire_routed_db(mock_prisma_client: Any, table: _SpendLogsTable) -> MagicMoc
return routed
async def _flush(mock_prisma_client: Any, logs: List[Dict[str, Any]]) -> None:
async def _flush(mock_prisma_client: MagicMock, logs: list[SpendLogRow]) -> None:
proxy_logging = MagicMock()
proxy_logging.failure_handler = AsyncMock()
await ProxyUpdateSpend.update_spend_logs(
@ -953,7 +962,7 @@ async def _flush(mock_prisma_client: Any, logs: List[Dict[str, Any]]) -> None:
@pytest.mark.asyncio
async def test_update_spend_logs_rekeys_the_rows_a_reused_provider_response_id_would_drop(
mock_prisma_client: Any, make_spend_log_row: Any
mock_prisma_client: MagicMock, make_spend_log_row: SpendLogRowFactory
) -> None:
"""Three requests answered with one provider id are three charged requests, so three rows
must land: the first keeps the provider id, the other two are re-keyed on their own call id.
@ -975,9 +984,28 @@ async def test_update_spend_logs_rekeys_the_rows_a_reused_provider_response_id_w
assert table.inserts == [["chatcmpl-static-1"] * 3, ["call-1", "call-2"]]
@pytest.mark.asyncio
@pytest.mark.parametrize(
"call_type",
["acompletion", "atext_completion", "aembedding", "aresponses", "aanthropic_messages", "allm_passthrough_route"],
)
async def test_update_spend_logs_rekeys_every_inference_call_type(
mock_prisma_client: MagicMock, make_spend_log_row: SpendLogRowFactory, call_type: str
) -> None:
"""A self-hosted server reuses its id on every inference route it serves, not only chat."""
table = _wire_spend_logs_table(mock_prisma_client)
logs = [
make_spend_log_row(request_id="static-1", litellm_call_id=f"call-{i}", call_type=call_type) for i in range(2)
]
await _flush(mock_prisma_client, logs)
assert list(table.rows) == ["static-1", "call-1"]
@pytest.mark.asyncio
async def test_update_spend_logs_reads_stored_identities_back_from_the_writer(
mock_prisma_client: Any, make_spend_log_row: Any
mock_prisma_client: MagicMock, make_spend_log_row: SpendLogRowFactory
) -> None:
"""The read-back must see the rows this flush just landed, so it goes to the writer: with a
read replica configured, ``db.query_raw`` is routed to the reader, and a lagging reader would
@ -998,7 +1026,7 @@ async def test_update_spend_logs_reads_stored_identities_back_from_the_writer(
@pytest.mark.asyncio
async def test_update_spend_logs_rekeys_only_the_colliding_rows_of_a_mixed_batch(
mock_prisma_client: Any, make_spend_log_row: Any
mock_prisma_client: MagicMock, make_spend_log_row: SpendLogRowFactory
) -> None:
"""Two providers reusing ids and one issuing unique ids share a flush: every row of the
unique provider keeps its key and only the later rows of each reused id are re-keyed."""
@ -1026,7 +1054,7 @@ async def test_update_spend_logs_rekeys_only_the_colliding_rows_of_a_mixed_batch
@pytest.mark.asyncio
async def test_update_spend_logs_rekeys_a_row_whose_provider_id_an_earlier_flush_stored(
mock_prisma_client: Any, make_spend_log_row: Any
mock_prisma_client: MagicMock, make_spend_log_row: SpendLogRowFactory
) -> None:
"""The colliding row usually landed in an earlier flush, or from another worker."""
table = _wire_spend_logs_table(
@ -1049,7 +1077,7 @@ async def test_update_spend_logs_rekeys_a_row_whose_provider_id_an_earlier_flush
@pytest.mark.asyncio
async def test_update_spend_logs_does_not_rekey_a_replay_of_a_stored_row(
mock_prisma_client: Any, make_spend_log_row: Any
mock_prisma_client: MagicMock, make_spend_log_row: SpendLogRowFactory
) -> None:
"""A transport retry replays the whole batch; a row already stored under its own call id is
the same request, not a collision, and must not become a second row."""
@ -1069,14 +1097,31 @@ async def test_update_spend_logs_does_not_rekey_a_replay_of_a_stored_row(
@pytest.mark.asyncio
@pytest.mark.parametrize(
"call_type", ["acreate_batch", "aretrieve_batch", "acreate_file", "avector_store_retrieve", "avector_store_delete"]
"call_type",
[
"acreate_batch",
"aretrieve_batch",
"acancel_batch",
"acreate_file",
"afile_retrieve",
"afile_content",
"afile_delete",
"aretrieve_fine_tuning_job",
"avideo_retrieve",
"aretrieve_container",
"avector_store_retrieve",
"avector_store_delete",
"aget_responses",
None,
],
)
async def test_update_spend_logs_keeps_object_keyed_rows_collapsed_on_their_object_id(
mock_prisma_client: Any, make_spend_log_row: Any, call_type: str
mock_prisma_client: MagicMock, make_spend_log_row: SpendLogRowFactory, call_type: str | None
) -> None:
"""Every poll of one batch shares its cost row by design, and every read of a stored object
is keyed on that object's id, so a duplicate here is not a lost row and the identity read-back
is not even issued."""
is not even issued. Only the inference call types are re-keyed; anything else, a call type
this code has never heard of included, keeps collapsing."""
table = _wire_spend_logs_table(mock_prisma_client)
logs = [
make_spend_log_row(request_id="batch_abc_batch_cost", litellm_call_id=f"call-{i}", call_type=call_type)
@ -1092,7 +1137,7 @@ async def test_update_spend_logs_keeps_object_keyed_rows_collapsed_on_their_obje
@pytest.mark.asyncio
async def test_update_spend_logs_drops_and_logs_a_row_already_keyed_on_its_call_id(
mock_prisma_client: Any, make_spend_log_row: Any, caplog: pytest.LogCaptureFixture
mock_prisma_client: MagicMock, make_spend_log_row: SpendLogRowFactory, caplog: pytest.LogCaptureFixture
) -> None:
"""A row whose key already is its call id (a client pinning x-litellm-call-id on a failure
row, say) has nothing to fall back to: it stays dropped, but loudly."""
@ -1118,7 +1163,7 @@ async def test_update_spend_logs_drops_and_logs_a_row_already_keyed_on_its_call_
@pytest.mark.asyncio
async def test_update_spend_logs_rekey_that_collides_again_stops(
mock_prisma_client: Any, make_spend_log_row: Any
mock_prisma_client: MagicMock, make_spend_log_row: SpendLogRowFactory
) -> None:
"""A client pinning one x-litellm-call-id against a provider reusing one response id
collides on both keys. The stored row already keyed on that call id reads as this row's
@ -1142,7 +1187,7 @@ async def test_update_spend_logs_rekey_that_collides_again_stops(
@pytest.mark.asyncio
async def test_update_spend_logs_treats_a_replay_of_a_rekeyed_row_as_stored(
mock_prisma_client: Any, make_spend_log_row: Any, caplog: pytest.LogCaptureFixture
mock_prisma_client: MagicMock, make_spend_log_row: SpendLogRowFactory, caplog: pytest.LogCaptureFixture
) -> None:
"""A transport retry replays a whole batch, including rows an earlier attempt already re-keyed.
Those are stored under their call id, so the read-back finds them there and nothing is re-keyed
@ -1166,7 +1211,7 @@ async def test_update_spend_logs_treats_a_replay_of_a_rekeyed_row_as_stored(
@pytest.mark.asyncio
async def test_update_spend_logs_leaves_rows_skipped_when_the_read_back_fails_on_its_data(
mock_prisma_client: Any, make_spend_log_row: Any, caplog: pytest.LogCaptureFixture
mock_prisma_client: MagicMock, make_spend_log_row: SpendLogRowFactory, caplog: pytest.LogCaptureFixture
) -> None:
"""A read-back the database rejects (not a transport fault) must not requeue the whole flush
behind it: the skipped rows stay skipped, which is what happened before, and the failure is
@ -1197,11 +1242,11 @@ async def test_update_spend_logs_leaves_rows_skipped_when_the_read_back_fails_on
@pytest.mark.asyncio
async def test_update_spend_logs_retries_the_flush_when_the_read_back_hits_a_transport_fault(
mock_prisma_client: Any, make_spend_log_row: Any
mock_prisma_client: MagicMock, make_spend_log_row: SpendLogRowFactory
) -> None:
"""A transport fault during the read-back is the same outage as one during the insert, so
the flush retries (and finally requeues) instead of leaving the rows behind."""
table = _wire_spend_logs_table(mock_prisma_client)
_ = _wire_spend_logs_table(mock_prisma_client)
mock_prisma_client.db.query_raw = AsyncMock(side_effect=httpx.ReadError("network blip"))
proxy_logging = MagicMock()
proxy_logging.failure_handler = AsyncMock()