test(e2e): guard destructive spend-log truncate behind explicit opt-in

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
yassin 2026-07-17 17:15:23 +00:00
parent adb1ffb119
commit b2c05fdd2c
3 changed files with 36 additions and 34 deletions

View file

@ -14,14 +14,14 @@ shared fixtures build on it.
"""
import functools
import sys
import os
from collections.abc import Generator, Iterator
from pathlib import Path
import pytest
import requests
from e2e_config import CONTROL_PLANE_BASE_URL, PROXY_BASE_URL
from e2e_db import reset_spend_logs
from e2e_result_reporter import covers_from_item, format_e2e_result_line, result_from_pytest
from lifecycle import GatewayProvider, ResourceManager
@ -112,25 +112,20 @@ def pytest_runtest_makereport(
def pytest_sessionfinish(session: pytest.Session, exitstatus: int) -> None:
"""Once the whole e2e session is done (all suites), truncate the spend logs so
the DB doesn't accumulate test rows. Sessions where no e2e test body ran leave
the DB alone so a `DATABASE_URL` pointing at a shared instance is never wiped
without an e2e run. Best-effort: a cleanup failure (no DB reachable) must not
fail the run. The spend_tracking dir goes on sys.path only for this import and
is removed after, so a broader `pytest tests/` run is not left with a mutated
path."""
the DB doesn't accumulate test rows. The truncate is destructive against
whatever `DATABASE_URL` points at, so it needs an explicit opt-in
(`E2E_RESET_SPEND_LOGS=1`): a local run against a shared or staging DB never
wipes real spend data unless the operator asked for it. On top of the opt-in
the DB is only touched when an e2e test body actually ran. Best-effort: a
cleanup failure (no DB reachable) must not fail the run."""
if os.environ.get("E2E_RESET_SPEND_LOGS") != "1":
return
if not session.stash.get(_E2E_TEST_RAN, False):
return
spend_dir = str(Path(__file__).parent / "quota_management" / "spend_tracking")
sys.path.insert(0, spend_dir)
try:
from spend_e2e_client import reset_spend_logs # pyright: ignore
reset_spend_logs()
except Exception as exc: # noqa: BLE001 - cleanup is best-effort
print(f"spend-log cleanup best-effort failed: {exc}")
finally:
if spend_dir in sys.path:
sys.path.remove(spend_dir)
try:
from bob_the_builder import remediate

26
tests/e2e/e2e_db.py Normal file
View file

@ -0,0 +1,26 @@
"""Shared e2e database helpers usable across suites without sys.path hacks.
Lives alongside e2e_config.py / lifecycle.py so any suite (or the top-level
conftest cleanup) can import it normally with `from e2e_db import ...`.
"""
from __future__ import annotations
import os
def reset_spend_logs() -> None:
"""Truncate LiteLLM_SpendLogs for a clean slate. No proxy endpoint deletes
spend logs (/global/spend/reset keeps them), so go to the DB directly. Uses
DATABASE_URL (default: the local docker postgres on its mapped host port; note
the in-container `@db` host isn't resolvable from the host, so default to
localhost).
"""
import psycopg
url = os.environ.get(
"DATABASE_URL",
"postgresql://llmproxy:dbpassword9090@localhost:5432/litellm",
)
with psycopg.connect(url) as conn:
_ = conn.execute('TRUNCATE TABLE "LiteLLM_SpendLogs"')

View file

@ -11,7 +11,6 @@ helpers from one place.
from __future__ import annotations
import os
import time
from collections.abc import Callable
from dataclasses import dataclass
@ -49,7 +48,6 @@ from models import (
__all__ = [
"SpendClient",
"build_client",
"reset_spend_logs",
"unique_marker",
"unwrap",
"is_ok",
@ -58,23 +56,6 @@ __all__ = [
]
def reset_spend_logs() -> None:
"""Truncate LiteLLM_SpendLogs for a clean slate. No proxy endpoint deletes
spend logs (/global/spend/reset keeps them), so go to the DB directly. Uses
DATABASE_URL (default: the local docker postgres on its mapped host port; note
the in-container `@db` host isn't resolvable from the host, so default to
localhost).
"""
import psycopg
url = os.environ.get(
"DATABASE_URL",
"postgresql://llmproxy:dbpassword9090@localhost:5432/litellm",
)
with psycopg.connect(url) as conn:
_ = conn.execute('TRUNCATE TABLE "LiteLLM_SpendLogs"')
def _chat_body(
model: str,
content: str,