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

test(integration): settle rollup and fallback row polling

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 00:52:25 +00:00
parent dc9889a481
commit 6e77d23f4d
3 changed files with 65 additions and 28 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

9
tests/integration/_support/client.py
View file

@ -34,12 +34,19 @@ def delete_key_if_present(candidate: Gateway, key: str) -> None:
assert read_rows('SELECT token FROM "LiteLLM_VerificationToken" WHERE token=%s', (digest,)) == []
def eventually(read: Callable[[], T], satisfied: Callable[[T], bool], seconds: float = 10) -> T:
def eventually(
read: Callable[[], T],
satisfied: Callable[[T], bool],
seconds: float = 10,
return_last_on_timeout: bool = False,
) -> T:
deadline: Final = time.monotonic() + seconds
while True:
observed: Final = read()
if satisfied(observed):
return observed
if return_last_on_timeout and time.monotonic() >= deadline:
return observed
assert time.monotonic() < deadline, f"State did not converge: {observed!r}"
time.sleep(0.1)

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

@ -123,22 +123,33 @@ def poll_cost_row(key: str) -> CostRow:
return result
def poll_rows(key: str, count: int) -> tuple[CostRow, ...]:
def read_rows_now(key: str) -> tuple[CostRow, ...]:
digest: Final = sha256(key.encode()).hexdigest()
rows: Final = read_rows(
'SELECT spend, status, metadata, prompt_tokens, completion_tokens, model_id '
'FROM "LiteLLM_SpendLogs" WHERE api_key=%s ORDER BY "startTime"',
(digest,),
)
return tuple(parsed for row in rows if (parsed := _row(row)) is not None)
def read() -> tuple[CostRow, ...]:
rows: Final = read_rows(
'SELECT spend, status, metadata, prompt_tokens, completion_tokens, model_id '
'FROM "LiteLLM_SpendLogs" WHERE api_key=%s ORDER BY "startTime"',
(digest,),
)
return tuple(parsed for row in rows if (parsed := _row(row)) is not None)
result: Final = eventually(read, lambda rows: len(rows) >= count, seconds=60)
def poll_rows(key: str, count: int) -> tuple[CostRow, ...]:
result: Final = eventually(
lambda: read_rows_now(key),
lambda rows: len(rows) >= count,
seconds=60,
)
return result
def poll_rollups(key: str, team_id: str, user_id: str, end_user_id: str, requests: int, spend: float) -> Rollups:
def poll_rollups(
key: str,
team_id: str,
user_id: str,
end_user_id: str,
target_spend: float,
target_requests: int,
) -> Rollups:
digest: Final = sha256(key.encode()).hexdigest()
def read() -> Rollups | None:
@ -178,21 +189,28 @@ def poll_rollups(key: str, team_id: str, user_id: str, end_user_id: str, request
daily_user=DailySpend.model_validate(daily_user_rows[0]),
daily_team=DailySpend.model_validate(daily_team_rows[0]),
)
if rollups.daily_user.api_requests < requests or rollups.daily_team.api_requests < requests:
return None
if not all(
approx_equal(actual, spend)
for actual in (
rollups.key_spend,
rollups.team_spend,
rollups.user_spend,
rollups.end_user_spend,
)
):
return None
return rollups
result: Final = eventually(read, lambda value: value is not None, seconds=60)
def settled(value: Rollups | None) -> bool:
return value is not None and all(
(
approx_equal(value.key_spend, target_spend),
approx_equal(value.team_spend, target_spend),
approx_equal(value.user_spend, target_spend),
approx_equal(value.end_user_spend, target_spend),
approx_equal(value.daily_user.spend, target_spend),
approx_equal(value.daily_team.spend, target_spend),
value.daily_user.api_requests == target_requests,
value.daily_team.api_requests == target_requests,
)
)
result: Final = eventually(
read,
settled,
seconds=20,
return_last_on_timeout=True,
)
assert result is not None
return result

20
tests/integration/cost_calculation/test_cost_tracking.py
View file

@ -7,6 +7,7 @@ from itertools import islice
import json
from hashlib import sha256
import struct
import time
from typing import Final, cast
import uuid
import wave
@ -27,6 +28,7 @@ from integration.cost_calculation.conftest import (
poll_failure_row,
poll_rollups,
poll_rows,
read_rows_now,
register_scenario_deployment,
)
from integration.cost_calculation.cost_tracking_case import (
@ -388,9 +390,11 @@ def test_case_bills_expected_cost(gateway: Gateway, case: CostTrackingTestCase)
assert isinstance(expected, ExactExpected)
if fallback_deployment is not None:
assert deployment is not None
assert len(rows) == 1
assert rows[0].status == "success"
assert rows[0].model_id == deployment.identity
time.sleep(3)
settled_rows: Final = read_rows_now(key)
assert len(settled_rows) == 1
assert settled_rows[0].status == "success"
assert settled_rows[0].model_id == deployment.identity
if isinstance(case.response, BinaryResponse):
header: Final = response.headers.get("x-litellm-response-cost")
if header is not None:
@ -413,7 +417,15 @@ def test_case_bills_expected_cost(gateway: Gateway, case: CostTrackingTestCase)
assert deployment is not None and team_id is not None and user_id is not None
assert end_user_id is not None
target_spend: Final = expected.spend * 3
rollups: Final = poll_rollups(key, team_id, user_id, end_user_id, requests=3, spend=target_spend)
target_requests: Final = 3
rollups: Final = poll_rollups(
key,
team_id,
user_id,
end_user_id,
target_spend,
target_requests,
)
assert approx_equal(rollups.key_spend, target_spend)
assert approx_equal(rollups.team_spend, target_spend)
assert approx_equal(rollups.user_spend, target_spend)