From f889b2b587d522944e2c5bd80e10112a890539eb Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Wed, 11 Nov 2026 12:36:02 -0800 Subject: [PATCH] test(integration): fix the baseline accounting and series cap flakes (#51096) * test(integration): run baseline accounting store tests in a private schema The accounting workers share their database with two gateway proxies, and every spend-log tick on those proxies projects any dirty baseline comparison it finds. When one of them published the test's scope between its last append and its own project() call, the test saw "unchanged" instead of "published". The module fixture now copies the twelve tables the store touches into a fresh schema and points DATABASE_URL at it for the module, so the gateway proxies, which only read public, can no longer publish these rows. * test(integration): wait for H9's own prometheus increments before it returns H9 stopped as soon as both workers' sample files already counted `other`, so its own chat's increment could still be in flight and land inside the next cell's window, which then saw two requests instead of one. H9 now takes a before scrape, counts the chats it sends, and finishes with the same _overflowed wait every other TestCapped cell uses. --- .../test_prometheus_series_cap.py | 23 ++++++----- .../spend/test_baseline_accounting.py | 38 ++++++++++++++++++- 2 files changed, 49 insertions(+), 12 deletions(-) diff --git a/tests/integration/observability/test_prometheus_series_cap.py b/tests/integration/observability/test_prometheus_series_cap.py index ecdb6dabfc5..68c96a257b9 100644 --- a/tests/integration/observability/test_prometheus_series_cap.py +++ b/tests/integration/observability/test_prometheus_series_cap.py @@ -7,6 +7,7 @@ from __future__ import annotations import time from collections.abc import Iterator, Sequence +from itertools import count from typing import Final import anthropic @@ -221,19 +222,21 @@ class TestCapped: """H9: fresh connections reach both workers, and each worker's own sample file names only the warm aliases while counting the fourth key on `other`, since the admitted sets live in the shared directory.""" extra: Final = capped.key("h9") + before: Final = scrape(capped.gateway) + sent: Final = count(1) - def send_on_a_fresh_connection() -> tuple[WorkerSamples, ...]: + def send_on_a_fresh_connection() -> tuple[int, tuple[WorkerSamples, ...]]: assert capped.chat(extra, Call.new()).status_code == 200 - return worker_samples(capped.prom_dir, REQUESTS) + return next(sent), worker_samples(capped.prom_dir, REQUESTS) - workers: Final = eventually( - send_on_a_fresh_connection, - lambda found: ( - sum(1 for worker in found if worker.overflow > 0) >= 2 - or any(extra.alias in worker.aliases for worker in found) - ), - seconds=90, - ) + def reached_both_workers(attempt: tuple[int, tuple[WorkerSamples, ...]]) -> bool: + _, found = attempt + return sum(1 for worker in found if worker.overflow > 0) >= 2 or any( + extra.alias in worker.aliases for worker in found + ) + + chats, workers = eventually(send_on_a_fresh_connection, reached_both_workers, seconds=90) + _overflowed(capped, extra, before, chats) assert all(extra.alias not in worker.aliases for worker in workers), workers assert sum(1 for worker in workers if worker.overflow > 0) >= 2, workers assert frozenset().union(*(worker.aliases for worker in workers)) == capped.warm_aliases, workers diff --git a/tests/integration/spend/test_baseline_accounting.py b/tests/integration/spend/test_baseline_accounting.py index 60c5c32511a..d384d509fd7 100644 --- a/tests/integration/spend/test_baseline_accounting.py +++ b/tests/integration/spend/test_baseline_accounting.py @@ -1,11 +1,13 @@ import asyncio import json +import os import uuid from collections.abc import AsyncIterator, Callable from contextlib import asynccontextmanager from dataclasses import replace from datetime import datetime, timedelta, timezone from typing import Final +from urllib.parse import parse_qsl, urlencode, urlsplit, urlunsplit import pytest import pytest_asyncio @@ -29,10 +31,42 @@ from tests.integration._support.prisma_clients import connected_prisma pytestmark = pytest.mark.asyncio(loop_scope="session") +_ISOLATED_TABLES: Final = ( + "LiteLLM_AutoRouterBaselineComparison", + "LiteLLM_AutoRouterBaselineObservation", + "LiteLLM_AutoRouterDailySpend", + "LiteLLM_AutoRouterSession", + "LiteLLM_AutoRouterUserSession", + "LiteLLM_DailyAgentSpend", + "LiteLLM_DailyEndUserSpend", + "LiteLLM_DailyOrganizationSpend", + "LiteLLM_DailyTagSpend", + "LiteLLM_DailyTeamSpend", + "LiteLLM_DailyUserSpend", + "LiteLLM_SpendLogs", +) + + +def _with_schema(database_url: str, schema: str) -> str: + parts: Final = urlsplit(database_url) + kept: Final = tuple((name, value) for name, value in parse_qsl(parts.query) if name != "schema") + return urlunsplit(parts._replace(query=urlencode((*kept, ("schema", schema))))) + + @pytest_asyncio.fixture(scope="module", loop_scope="session") async def db() -> AsyncIterator[Prisma]: - async with connected_prisma() as client: - yield client + schema: Final = f"baseline_accounting_{uuid.uuid4().hex}" + async with connected_prisma() as admin: + await admin.execute_raw(f'CREATE SCHEMA "{schema}"') + try: + for table in _ISOLATED_TABLES: + await admin.execute_raw(f'CREATE TABLE "{schema}"."{table}" (LIKE "{table}" INCLUDING ALL)') + with pytest.MonkeyPatch.context() as environment: + environment.setenv("DATABASE_URL", _with_schema(os.environ["DATABASE_URL"], schema)) + async with connected_prisma() as client: + yield client + finally: + await admin.execute_raw(f'DROP SCHEMA "{schema}" CASCADE') @asynccontextmanager