scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-11 03:38:38 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 84

test(integration): wait for the batch retrieval row and drop the failed-only output file

Browse source 
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
kerry 2026-09-20 01:46:23 +00:00
parent 2e16cd76c1
commit 8d6326356d
2 changed files with 55 additions and 39 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

12
tests/integration/cost_calculation/conftest.py
View file

@ -3,7 +3,7 @@ from __future__ import annotations
import functools
import json
import os
from collections.abc import Mapping
from collections.abc import Callable, Mapping
from dataclasses import dataclass
from hashlib import sha256
from typing import Final
@ -134,8 +134,16 @@ def read_rows_now(key: str) -> tuple[CostRow, ...]:
def poll_rows(key: str, count: int) -> tuple[CostRow, ...]:
return poll_rows_where(key, count, lambda _row: True)
def poll_rows_where(
key: str,
count: int,
predicate: Callable[[CostRow], bool],
) -> tuple[CostRow, ...]:
result: Final = eventually(
lambda: read_rows_now(key),
lambda: tuple(row for row in read_rows_now(key) if predicate(row)),
lambda rows: len(rows) >= count,
seconds=60,
)

82
tests/integration/cost_calculation/test_batch_realtime_cost.py
View file

@ -12,7 +12,7 @@ import websockets
from integration._support.client import JSON_OBJECT, Gateway, Scenario, object_value, string_value
from integration._support.upstream import delete_scenario, register_scenario
from integration.cost_calculation.assertions import assert_exact
from integration.cost_calculation.conftest import poll_rows, read_rows_now
from integration.cost_calculation.conftest import poll_rows, poll_rows_where, read_rows_now
from integration.cost_calculation.cost_tracking_case import (
BATCH_CASES,
REALTIME_CASES,
@ -67,6 +67,8 @@ def _batch_response(case: BatchCostCase) -> JsonResponse | RoutedResponse:
"completed": case.completed_count,
"failed": case.failed_count,
}
has_output: Final = any(line.status_code == 200 for line in case.output_lines)
has_failed: Final = any(line.status_code != 200 for line in case.output_lines)
batch: Final = {
"id": "batch-$REQUEST_ID",
"object": "batch",
@ -75,8 +77,8 @@ def _batch_response(case: BatchCostCase) -> JsonResponse | RoutedResponse:
"input_file_id": "file-in-$REQUEST_ID",
"completion_window": "24h",
"status": "completed",
"output_file_id": "file-out-$REQUEST_ID" if case.output_lines else None,
"error_file_id": None if case.output_lines else "file-err-$REQUEST_ID",
"output_file_id": "file-out-$REQUEST_ID" if has_output else None,
"error_file_id": "file-err-$REQUEST_ID" if has_failed else None,
"created_at": 1,
"in_progress_at": 1,
"completed_at": 1,
@ -84,39 +86,46 @@ def _batch_response(case: BatchCostCase) -> JsonResponse | RoutedResponse:
"request_counts": counts,
"metadata": None,
}
routes: Final = {
"POST /files": JsonResponse(
content_type="application/json",
body={
"id": "file-in-$REQUEST_ID",
"object": "file",
"purpose": "batch",
"bytes": 100,
"created_at": 1,
"filename": "in.jsonl",
"status": "processed",
},
),
"POST /batches": JsonResponse(
content_type="application/json",
body={
**batch,
"status": "validating",
"output_file_id": None,
"error_file_id": None,
},
),
"GET /batches/batch-$REQUEST_ID": JsonResponse(
content_type="application/json",
body=batch,
),
**(
{
"GET /files/file-out-$REQUEST_ID/content": TextResponse(
content_type="application/jsonl",
body="\n".join(lines) + ("\n" if lines else ""),
)
}
if has_output
else {}
),
}
return RoutedResponse(
content_type="application/x-routed",
routes={
"POST /files": JsonResponse(
content_type="application/json",
body={
"id": "file-in-$REQUEST_ID",
"object": "file",
"purpose": "batch",
"bytes": 100,
"created_at": 1,
"filename": "in.jsonl",
"status": "processed",
},
),
"POST /batches": JsonResponse(
content_type="application/json",
body={
**batch,
"status": "validating",
"output_file_id": None,
"error_file_id": None,
},
),
"GET /batches/batch-$REQUEST_ID": JsonResponse(
content_type="application/json",
body=batch,
),
"GET /files/file-out-$REQUEST_ID/content": TextResponse(
content_type="application/jsonl",
body="\n".join(lines) + ("\n" if lines else ""),
),
},
routes=routes,
)
@ -164,7 +173,6 @@ def test_batch_costs(gateway: Gateway, case: BatchCostCase) -> None:
)
assert file_response.is_success, file_response.text
file_body: Final = JSON_OBJECT.validate_json(file_response.content)
time.sleep(2)
batch_response: Final = gateway.request(
"POST",
"/v1/batches",
@ -183,9 +191,9 @@ def test_batch_costs(gateway: Gateway, case: BatchCostCase) -> None:
second_retrieval: Final = gateway.request("GET", f"/v1/batches/{batch_id}", key=key)
assert first_retrieval.is_success, first_retrieval.text
assert second_retrieval.is_success, second_retrieval.text
rows: Final = poll_rows(key, 1)
retrieval_rows: Final = tuple(row for row in rows if row.call_type == "aretrieve_batch")
retrieval_rows: Final = poll_rows_where(key, 1, lambda row: row.call_type == "aretrieve_batch")
assert len(retrieval_rows) == 1
rows: Final = read_rows_now(key)
assert all(row.spend == 0.0 for row in rows if row.call_type != "aretrieve_batch")
row: Final = retrieval_rows[0]
assert row.status == "success"