scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-04 02:31:27 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 71

test(batches): add failing regression for completed-batch cost tracking gap

Browse source 
Adds a regression test demonstrating that a user GET that transitions a batch to
complete marks batch_processed=True on a plain status sync. Because CheckBatchCost
only picks up batches with batch_processed=False, any batch whose asynchronous cost
callback failed is permanently skipped and its cost is never recorded, causing
billing discrepancies.

The test encodes the expected behavior (a status sync must not mark a batch
cost-tracked) and currently fails against the existing code, pinpointing the bug to
fix. It will pass once batch_processed is set only from the cost-recording path.
This commit is contained in:
Claude 2026-06-08 19:36:58 +00:00 • committed by mubashir1osmani
parent e6db497fb8
commit eef5bd1ad2
No known key found for this signature in database
GPG key ID: AB055FF67D0B4D9A
1 changed files with 51 additions and 0 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

51
tests/test_litellm/enterprise/proxy/test_batch_update_db_managed_output_file_id.py
View file

@ -52,6 +52,57 @@ def _build_prisma_mock():
return mock
@pytest.mark.asyncio
async def test_status_sync_to_complete_does_not_mark_batch_cost_tracked():
"""
Regression for the batch cost-tracking gap.
When a user GET /v1/batches/{id} transitions a batch from validating ->
completed, the endpoint calls update_batch_in_database to sync the status. The
cost callback for that retrieve is enqueued asynchronously and can fail. The
batch_processed flag means "cost has been recorded" and is what CheckBatchCost
uses (batch_processed=False) to find batches whose cost still needs tracking.
If update_batch_in_database marks batch_processed=True just for syncing status,
a batch whose async cost callback failed is permanently hidden from CheckBatchCost
and its cost is never recorded. So a plain status sync to complete must NOT set
batch_processed=True; that flag belongs to the cost-recording path. This test
fails while the bug is present and passes once the flag is no longer set here.
"""
batch_id = "batch_complete_status_sync"
unified_batch_id = (
"litellm_proxy;model_id:m;llm_batch_id:batch_complete_status_sync"
)
response = _build_batch_response(
batch_id=batch_id,
status="completed",
output_file_id=None,
hidden_params={"model_id": "m", "model_name": "openai/gpt-4o"},
)
mock_prisma = _build_prisma_mock()
await update_batch_in_database(
batch_id=batch_id,
unified_batch_id=unified_batch_id,
response=response,
managed_files_obj=_build_managed_files_mock(),
prisma_client=mock_prisma,
verbose_proxy_logger=MagicMock(),
user_api_key_dict=UserAPIKeyAuth(user_id="user-abc"),
)
update_data = mock_prisma.db.litellm_managedobjecttable.update.call_args.kwargs[
"data"
]
assert update_data["status"] == "complete"
assert update_data.get("batch_processed") is not True, (
"A status sync to complete must not mark batch_processed=True; doing so hides "
"batches whose async cost callback failed from CheckBatchCost, so their cost "
"is never recorded."
)
@pytest.mark.asyncio
async def test_update_batch_in_database_stores_unified_output_file_id():
raw_output_file_id = "file-rawoutput789"