From 4449c3cfd30227cb2c94a194d0f16bd2e5b836ee Mon Sep 17 00:00:00 2001 From: kerry Date: Tue, 22 Sep 2026 20:50:33 +0000 Subject: [PATCH] test(integration): failed daily user rollup commits are retried so spend report and daily activity agree (Pylon #7268) Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- tests/integration/contracts.json | 3 + .../spend/test_daily_rollup_retry.py | 182 ++++++++++++++++++ 2 files changed, 185 insertions(+) create mode 100644 tests/integration/spend/test_daily_rollup_retry.py diff --git a/tests/integration/contracts.json b/tests/integration/contracts.json index ff9ef2b3903..20cb34a5031 100644 --- a/tests/integration/contracts.json +++ b/tests/integration/contracts.json @@ -284,6 +284,9 @@ "tests/integration/compatibility/test_openai_consumer.py::test_retained_openai_clients_parse_real_proxy_tool_and_usage_responses": [ "other.compatibility.openai.retained_client_parses_tools_and_usage" ], + "tests/integration/spend/test_daily_rollup_retry.py::test_failed_daily_user_rollup_commit_is_retried_so_spend_report_and_daily_activity_agree": [ + "spend.daily_rollup.failed_user_commit_is_retried_until_report_and_daily_activity_agree" + ], "tests/integration/spend/test_filtered_ledger.py::test_rotated_keys_users_and_model_groups_preserve_success_failure_cache_ledger": [ "quota_management.spend_tracking.filtered_ledger_preserves_owner_identity_and_totals" ], diff --git a/tests/integration/spend/test_daily_rollup_retry.py b/tests/integration/spend/test_daily_rollup_retry.py new file mode 100644 index 00000000000..cf1b989639a --- /dev/null +++ b/tests/integration/spend/test_daily_rollup_retry.py @@ -0,0 +1,182 @@ +import json +import os +import uuid +from collections.abc import Iterable +from hashlib import sha256 +from typing import Final + +import psycopg +import pytest +from integration._support.client import Gateway, delete_key_if_present, eventually, string_value +from integration._support.database import read_rows +from integration._support.wire import Reply, Request, wire_server +from psycopg import sql + + +def _execute(statements: Iterable[sql.Composable]) -> None: + with psycopg.connect(os.environ["DATABASE_URL"], autocommit=True) as connection: + for statement in statements: + connection.execute(statement) + + +def _install_daily_user_rollup_fault(user_id: str) -> str: + suffix: Final = f"fault-{uuid.uuid4().hex}" + sequence: Final = sql.Identifier(f"{suffix}_attempts") + function: Final = sql.Identifier(suffix) + _execute( + ( + sql.SQL("CREATE SEQUENCE {}").format(sequence), + sql.SQL( + "CREATE FUNCTION {}() RETURNS trigger LANGUAGE plpgsql AS $fault$ " + "BEGIN PERFORM nextval({}); " + "RAISE EXCEPTION 'synthetic daily rollup outage' USING ERRCODE = '55P03'; " + "END $fault$" + ).format(function, sql.Literal(f"{suffix}_attempts")), + sql.SQL( + 'CREATE TRIGGER {} BEFORE INSERT ON "LiteLLM_DailyUserSpend" ' + "FOR EACH ROW WHEN (NEW.user_id = {}) EXECUTE FUNCTION {}()" + ).format(sql.Identifier(suffix), sql.Literal(user_id), function), + ) + ) + return suffix + + +def _lift_daily_user_rollup_fault(suffix: str) -> None: + _execute( + ( + sql.SQL('DROP TRIGGER IF EXISTS {} ON "LiteLLM_DailyUserSpend"').format(sql.Identifier(suffix)), + sql.SQL("DROP FUNCTION IF EXISTS {}()").format(sql.Identifier(suffix)), + sql.SQL("DROP SEQUENCE IF EXISTS {}").format(sql.Identifier(f"{suffix}_attempts")), + ) + ) + + +def _rollup_attempts(suffix: str) -> int: + with psycopg.connect(os.environ["DATABASE_URL"], autocommit=True) as connection: + row: Final = connection.execute( + sql.SQL("SELECT CASE WHEN is_called THEN last_value ELSE 0 END FROM {}").format( + sql.Identifier(f"{suffix}_attempts") + ) + ).fetchone() + assert row is not None + return int(row[0]) + + +@pytest.mark.covers("spend.daily_rollup.failed_user_commit_is_retried_until_report_and_daily_activity_agree") +def test_failed_daily_user_rollup_commit_is_retried_so_spend_report_and_daily_activity_agree( + gateway: Gateway, +) -> None: + def provider(request: Request) -> Reply: + assert request.method == "POST" and request.target == "/v1/chat/completions" + return Reply( + body=json.dumps( + { + "id": "chatcmpl-" + uuid.uuid4().hex, + "object": "chat.completion", + "created": 1, + "model": "gpt-4o-mini", + "choices": [ + { + "index": 0, + "message": {"role": "assistant", "content": "synthetic rollup answer"}, + "finish_reason": "stop", + } + ], + "usage": {"prompt_tokens": 20, "completion_tokens": 20, "total_tokens": 40}, + } + ).encode() + ) + + with wire_server(provider) as wire, gateway.scenario() as scenario: + model: Final = scenario.model( + api_base=wire.url + "/v1", input_cost_per_token=0.001, output_cost_per_token=0.002, num_retries=0 + ) + user: Final = scenario.user() + key: Final = string_value(gateway.post("/key/generate", {"user_id": user, "models": [model]})["key"]) + scenario.cleanups.callback(delete_key_if_present, gateway, key) + digest: Final = sha256(key.encode()).hexdigest() + suffix: Final = _install_daily_user_rollup_fault(user) + scenario.cleanups.callback(_lift_daily_user_rollup_fault, suffix) + response: Final = gateway.request( + "POST", + "/v1/chat/completions", + {"model": model, "messages": [{"role": "user", "content": "rollup retry control"}]}, + key=key, + ) + assert response.status_code == 200, response.text + assert float(response.headers["x-litellm-response-cost"]) == pytest.approx(20 * 0.001 + 20 * 0.002) + body: Final = response.json() + spend_rows: Final = eventually( + lambda: read_rows( + 'SELECT spend, DATE("startTime")::text AS day, model FROM "LiteLLM_SpendLogs" WHERE request_id=%s', + (body["id"],), + ), + lambda values: len(values) == 1, + seconds=70, + ) + assert float(spend_rows[0]["spend"]) == pytest.approx(0.06) + day: Final = string_value(spend_rows[0]["day"]) + stored_model: Final = string_value(spend_rows[0]["model"]) + eventually(lambda: _rollup_attempts(suffix), lambda attempts: attempts >= 1, seconds=70) + _lift_daily_user_rollup_fault(suffix) + activity: Final = eventually( + lambda: gateway.request( + "GET", + "/user/daily/activity/aggregated", + params={"start_date": day, "end_date": day, "api_key": digest}, + ), + lambda polled: ( + polled.status_code == 200 + and len(polled.json().get("results", ())) > 0 + and polled.json()["results"][0]["breakdown"]["api_keys"] + .get(digest, {}) + .get("metrics", {}) + .get("spend", 0) + == pytest.approx(0.06) + ), + seconds=90, + ) + assert activity.status_code == 200, activity.text + metrics: Final = activity.json()["results"][0]["breakdown"]["api_keys"][digest]["metrics"] + assert metrics == { + "spend": pytest.approx(0.06), + "flat_cost": pytest.approx(0.0), + "prompt_tokens": 20, + "completion_tokens": 20, + "cache_read_input_tokens": 0, + "cache_creation_input_tokens": 0, + "compression_saved_tokens": 0, + "compression_savings_spend": pytest.approx(0.0), + "prompt_caching_savings_spend": pytest.approx(0.0), + "gateway_injected_caching_savings_spend": pytest.approx(0.0), + "autorouter_savings_spend": pytest.approx(0.0), + "total_tokens": 40, + "successful_requests": 1, + "failed_requests": 0, + "api_requests": 1, + "total_response_time_ms": metrics["total_response_time_ms"], + "timed_requests": metrics["timed_requests"], + } + report: Final = gateway.request( + "GET", + "/global/spend/report", + params={"start_date": day, "end_date": day, "api_key": digest}, + ) + assert report.status_code == 200, report.text + assert report.json() == [ + { + "api_key": digest, + "total_cost": pytest.approx(0.06), + "total_input_tokens": 20, + "total_output_tokens": 20, + "model_details": [ + { + "model": stored_model, + "total_cost": pytest.approx(0.06), + "total_input_tokens": 20, + "total_output_tokens": 20, + } + ], + } + ] + assert metrics["spend"] == pytest.approx(report.json()[0]["total_cost"])