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