From b2c05fdd2c1b35994e5ebd438a79bf79880c9529 Mon Sep 17 00:00:00 2001
From: yassin <yassin@berri.ai>
Date: Fri, 17 Jul 2026 17:15:23 +0000
Subject: [PATCH] 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>
---
 tests/e2e/conftest.py                         | 25 +++++++-----------
 tests/e2e/e2e_db.py                           | 26 +++++++++++++++++++
 .../spend_tracking/spend_e2e_client.py        | 19 --------------
 3 files changed, 36 insertions(+), 34 deletions(-)
 create mode 100644 tests/e2e/e2e_db.py

diff --git a/tests/e2e/conftest.py b/tests/e2e/conftest.py
index 3aec104c861..74a6e39e5d6 100644
--- a/tests/e2e/conftest.py
+++ b/tests/e2e/conftest.py
@@ -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
diff --git a/tests/e2e/e2e_db.py b/tests/e2e/e2e_db.py
new file mode 100644
index 00000000000..99764cf20ef
--- /dev/null
+++ b/tests/e2e/e2e_db.py
@@ -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"')
diff --git a/tests/e2e/quota_management/spend_tracking/spend_e2e_client.py b/tests/e2e/quota_management/spend_tracking/spend_e2e_client.py
index c4991199187..9b3a09d7dca 100644
--- a/tests/e2e/quota_management/spend_tracking/spend_e2e_client.py
+++ b/tests/e2e/quota_management/spend_tracking/spend_e2e_client.py
@@ -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,