scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-13 23:11:40 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 100

test(proxy): pin the row budget in the index fan-out tests

Browse source 
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
yucheng 2026-09-10 08:06:00 +00:00
parent 16a5d4df8d
commit e0eb761e2a
2 changed files with 8 additions and 7 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

6
tests/test_litellm/proxy/db/test_spend_log_tool_index.py
View file

@ -11,7 +11,7 @@ from unittest.mock import AsyncMock, MagicMock
import httpx
import pytest
from litellm.constants import SPEND_LOG_WRITE_BATCH_MAX_ROWS
from litellm.proxy.db import spend_log_tool_index
from litellm.proxy.db.spend_log_tool_index import (
ToolUsageTransaction,
build_tool_usage_transaction,
@ -281,7 +281,8 @@ class TestFlushToolUsageTransactions:
assert data["update"]["request_count"] == {"increment": 2}
@pytest.mark.asyncio
async def test_index_rows_are_written_in_bounded_statements_outside_the_rollup_transaction(self):
async def test_index_rows_are_written_in_bounded_statements_outside_the_rollup_transaction(self, monkeypatch):
monkeypatch.setattr(spend_log_tool_index, "SPEND_LOG_WRITE_BATCH_MAX_ROWS", 100)
prisma, batcher = _prisma_with_batcher()
tool_names = tuple(f"tool_{i}" for i in range(50))
transactions = [_transaction(f"r{i}", tool_names=tool_names) for i in range(5)]
@ -289,7 +290,6 @@ class TestFlushToolUsageTransactions:
statements = prisma.db.litellm_spendlogtoolindex.create_many.call_args_list
assert [len(call.kwargs["data"]) for call in statements] == [100, 100, 50]
assert all(len(call.kwargs["data"]) <= SPEND_LOG_WRITE_BATCH_MAX_ROWS for call in statements)
assert all(call.kwargs["skip_duplicates"] is True for call in statements)
assert _index_rows_written(prisma) == [
(txn.request_id, tool_name) for txn in transactions for tool_name in tool_names

9
tests/test_litellm/proxy/guardrails/test_usage_tracking.py
View file

@ -6,7 +6,7 @@ from unittest.mock import AsyncMock, MagicMock
import httpx
import pytest
from litellm.constants import SPEND_LOG_WRITE_BATCH_MAX_ROWS
from litellm.proxy.guardrails import usage_tracking
from litellm.proxy.guardrails.usage_tracking import (
_MAX_PENDING_ROWS,
PendingRollups,
@ -535,11 +535,12 @@ def _index_rows_written(prisma: MagicMock) -> list[tuple[str, str]]:
@pytest.mark.asyncio
async def test_index_rows_are_written_in_row_bounded_statements():
async def test_index_rows_are_written_in_row_bounded_statements(monkeypatch):
"""
LIT-5931: the drain caps logs, not logs x guardrails, so a fan-out must be
split into statements the query engine can afford instead of one create_many.
"""
monkeypatch.setattr(usage_tracking, "SPEND_LOG_WRITE_BATCH_MAX_ROWS", 100)
prisma = _prisma()
guardrail_ids = tuple(f"guard-{i}" for i in range(50))
logs = [_fan_out_payload(f"r{i}", guardrail_ids) for i in range(5)]
@ -548,13 +549,13 @@ async def test_index_rows_are_written_in_row_bounded_statements():
statements = prisma.db.litellm_spendlogguardrailindex.create_many.call_args_list
assert [len(call.kwargs["data"]) for call in statements] == [100, 100, 50]
assert all(len(call.kwargs["data"]) <= SPEND_LOG_WRITE_BATCH_MAX_ROWS for call in statements)
assert all(call.kwargs["skip_duplicates"] is True for call in statements)
assert _index_rows_written(prisma) == [(f"r{i}", gid) for i in range(5) for gid in guardrail_ids]
@pytest.mark.asyncio
async def test_one_failing_index_statement_does_not_drop_the_others_or_the_rollup():
async def test_one_failing_index_statement_does_not_drop_the_others_or_the_rollup(monkeypatch):
monkeypatch.setattr(usage_tracking, "SPEND_LOG_WRITE_BATCH_MAX_ROWS", 100)
prisma = _prisma()
prisma.db.litellm_spendlogguardrailindex.create_many.side_effect = [None, httpx.ReadTimeout("ambiguous"), None]
guardrail_ids = tuple(f"guard-{i}" for i in range(50))