From 6e77d23f4d50503da2b4d35ab883df89744e14f8 Mon Sep 17 00:00:00 2001 From: kerry Date: Sun, 20 Sep 2026 00:52:25 +0000 Subject: [PATCH] test(integration): settle rollup and fallback row polling Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- tests/integration/_support/client.py | 9 ++- .../integration/cost_calculation/conftest.py | 64 ++++++++++++------- .../cost_calculation/test_cost_tracking.py | 20 ++++-- 3 files changed, 65 insertions(+), 28 deletions(-) diff --git a/tests/integration/_support/client.py b/tests/integration/_support/client.py index 0b6771623c0..e07cbe6b2a3 100644 --- a/tests/integration/_support/client.py +++ b/tests/integration/_support/client.py @@ -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) diff --git a/tests/integration/cost_calculation/conftest.py b/tests/integration/cost_calculation/conftest.py index fb20f5a9cc2..bc68419f01d 100644 --- a/tests/integration/cost_calculation/conftest.py +++ b/tests/integration/cost_calculation/conftest.py @@ -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 diff --git a/tests/integration/cost_calculation/test_cost_tracking.py b/tests/integration/cost_calculation/test_cost_tracking.py index 692b76fff15..05346396a8a 100644 --- a/tests/integration/cost_calculation/test_cost_tracking.py +++ b/tests/integration/cost_calculation/test_cost_tracking.py @@ -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)