mirror of
https://github.com/BerriAI/litellm.git
synced 2026-08-28 05:25:59 +00:00
8d96e959db
tests/e2e/conftest.py's pytest_sessionfinish truncated LiteLLM_SpendLogs against whatever DATABASE_URL resolved to, gated only by "an e2e test body ran". Pointed at a shared or staging DB, a routine local run wiped real spend data. It also reached the truncate helper through a sys.path.insert into quota_management/spend_tracking/spend_e2e_client.py, a cross-suite import-by-path hack it then unwound in a finally. The cleanup now routes through a new run_spend_log_cleanup in a top-level tests/e2e/e2e_db.py, which fires the destructive truncate only when the operator set E2E_RESET_SPEND_LOGS=1 and an e2e test actually ran. Any other value (unset, 0, true, empty) leaves the DB untouched, so presence of the variable alone or a test run alone never arms the truncate. The decision plus the injectable truncate callable live in that pure helper, and conftest is a thin adapter that supplies os.environ.get(...), the session stash, and reset_spend_logs. reset_spend_logs itself moved from spend_e2e_client.py into e2e_db.py (implementation unchanged), sitting next to e2e_config and lifecycle so both conftest and any suite import it by name; the sys.path munging is gone. Nothing else imported reset_spend_logs, so spend_e2e_client.py drops the definition, its __all__ entry, and the now-unused os import.
56 lines
2.3 KiB
Python
56 lines
2.3 KiB
Python
"""Shared, destructive DB helpers for the e2e harness.
|
|
|
|
Kept at the top level next to e2e_config and lifecycle so every suite imports it
|
|
by name (`from e2e_db import ...`); no suite reaches into another's directory by
|
|
mutating sys.path.
|
|
|
|
reset_spend_logs truncates LiteLLM_SpendLogs and cannot be undone, so the
|
|
session-finish cleanup routes through run_spend_log_cleanup, which fires the
|
|
truncate only on an explicit operator opt-in. "An e2e test ran" is necessary but
|
|
never sufficient: a DATABASE_URL pointing at a shared or staging instance must
|
|
not be wiped by a routine local run that merely exercised a test.
|
|
"""
|
|
|
|
import os
|
|
from collections.abc import Callable
|
|
|
|
RESET_OPT_IN_ENV = "E2E_RESET_SPEND_LOGS"
|
|
|
|
|
|
def run_spend_log_cleanup(
|
|
*, opt_in: str | None, e2e_test_ran: bool, truncate: Callable[[], None]
|
|
) -> bool:
|
|
"""Invoke `truncate` iff the destructive spend-log reset is both opted into
|
|
and warranted, returning whether the truncate was attempted.
|
|
|
|
The truncate fires only when the opt-in value is exactly "1" AND an e2e test
|
|
body actually ran. Any other opt-in value (unset, "0", "true", "") leaves the
|
|
DB untouched, so the destructive path is never armed by the env var's mere
|
|
presence or by a test run on its own. Best-effort: a truncate failure is
|
|
swallowed so cleanup never fails the session, so the returned bool reports
|
|
that the reset was attempted, not that the DB call succeeded.
|
|
"""
|
|
if opt_in != "1" or not e2e_test_ran:
|
|
return False
|
|
try:
|
|
truncate()
|
|
except Exception as exc: # noqa: BLE001 - cleanup is best-effort
|
|
print(f"spend-log cleanup best-effort failed: {exc}")
|
|
return True
|
|
|
|
|
|
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; 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"')
|