mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-14 23:21:35 +00:00
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:
parent
adb1ffb119
commit
b2c05fdd2c
3 changed files with 36 additions and 34 deletions
|
|
@ -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
26
tests/e2e/e2e_db.py
Normal 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"')
|
||||
|
|
@ -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,
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue