scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-07 08:26:10 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 36

test(proxy/db): scrub DATABASE_URL in test_routing_prisma_wrapper to stop env leak

Browse source 
PrismaWrapper.get_rds_iam_token writes DATABASE_URL directly into
os.environ. monkeypatch.delenv tracks nothing when the variable was
unset to begin with, so the synthesized writer URL (e.g. postgresql://
litellm:WRITER-TOKEN@writer.aurora.local:5432/litellm) survives the
test and leaks across the xdist worker.

Downstream effect: tests/test_litellm/proxy/common_utils/
test_key_rotation_e2e.py::TestDeprecatedKeyLookupDbE2E::
test_deprecated_key_grace_period_cache_hit_path stops skipping
(os.getenv("DATABASE_URL") is now truthy) and tries to open a real
Prisma connection against the bogus host, failing with
httpx.ConnectError under proxy-infra CI.

Add an autouse snapshot/restore fixture for the managed DATABASE_*
env vars, mirroring the pattern already used in test_db_url_settings.py
where the same docstring already documents that monkeypatch cannot
undo direct os.environ mutations.

Co-authored-by: Krrish Dholakia <krrish-berri-2@users.noreply.github.com>
This commit is contained in:
Cursor Agent 2026-06-13 03:20:43 +00:00
parent 0dc203bd65
commit edd9f25d21
No known key found for this signature in database
1 changed files with 46 additions and 0 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

46
tests/test_litellm/proxy/db/test_routing_prisma_wrapper.py
View file

@ -18,6 +18,52 @@ sys.path.insert(0, os.path.abspath("../../../.."))
# stub the prisma SDK do so per-test via monkeypatch, which is properly scoped.
_MANAGED_DB_ENV_VARS = (
"IAM_TOKEN_DB_AUTH",
"DATABASE_URL",
"DATABASE_URL_READ_REPLICA",
"DATABASE_HOST",
"DATABASE_PORT",
"DATABASE_USER",
"DATABASE_USERNAME",
"DATABASE_NAME",
"DATABASE_SCHEMA",
"DATABASE_PASSWORD",
"DATABASE_HOST_READ_REPLICA",
"DATABASE_PORT_READ_REPLICA",
"DATABASE_USER_READ_REPLICA",
"DATABASE_USERNAME_READ_REPLICA",
"DATABASE_NAME_READ_REPLICA",
"DATABASE_SCHEMA_READ_REPLICA",
"DATABASE_PASSWORD_READ_REPLICA",
)
@pytest.fixture(autouse=True)
def _scrub_db_env():
"""Snapshot and restore DB env vars across every test in this module.
`PrismaWrapper.get_rds_iam_token` (and a few sibling code paths) write
`DATABASE_URL` directly into `os.environ`, which `monkeypatch` cannot
undo when the variable started out unset. Without this fixture the
synthesized writer URL leaks into later tests in the same xdist worker;
e.g. `tests/test_litellm/proxy/common_utils/test_key_rotation_e2e.py::
TestDeprecatedKeyLookupDbE2E::test_deprecated_key_grace_period_cache_hit_path`
stops skipping (`os.getenv("DATABASE_URL")` is now truthy) and tries to
open a Prisma connection against the bogus host, failing with
`httpx.ConnectError`. Mirrors the fixture in `test_db_url_settings.py`.
"""
saved = {var: os.environ.get(var) for var in _MANAGED_DB_ENV_VARS}
try:
yield
finally:
for var, value in saved.items():
if value is None:
os.environ.pop(var, None)
else:
os.environ[var] = value
def _make_wrappers():
from litellm.proxy.db.prisma_client import PrismaWrapper