diff --git a/tests/proxy_behavior/management/conftest.py b/tests/proxy_behavior/management/conftest.py index c16adc39db1..23f2e24e604 100644 --- a/tests/proxy_behavior/management/conftest.py +++ b/tests/proxy_behavior/management/conftest.py @@ -7,6 +7,8 @@ run. Tests assert at the HTTP boundary. import os import tempfile +import uuid +from dataclasses import dataclass from typing import AsyncIterator import httpx @@ -80,9 +82,66 @@ async def world(prisma): """The immutable read-world seed. Re-seeds at session start so each pytest invocation gets a clean world. - Tests must not mutate these rows; write tests use Slice 5's namespace + - teardown fixtures for scratch entities. + Tests must not mutate these rows; write tests use the ``scratch`` fixture + below for scoped entities that get torn down per-test. """ from .actors import seed_world return await seed_world(prisma) + + +SCRATCH_PREFIX = "scratch-" + + +@dataclass(frozen=True) +class Scratch: + """Per-test namespace for write scenarios. + + Tests must tag any entity they create with ``scratch.prefix`` in a column + the teardown filter inspects (``key_alias``, ``key_name``, ``team_alias``, + ``team_id``, ``user_id``, or ``budget_id``). Anything not tagged will be + left behind and pollute the next session. + """ + + prefix: str + + def tag(self, suffix: str = "") -> str: + return f"{self.prefix}-{suffix}" if suffix else self.prefix + + +@pytest_asyncio.fixture +async def scratch(prisma): + """Function-scoped scratch namespace + targeted delete_many teardown. + + The teardown deletes any rows on the volatile tables whose namespace column + starts with ``scratch.prefix``. Per CLAUDE.md, this is Prisma-only — no raw + SQL — and uses ``delete_many`` to batch the writes. + """ + handle = Scratch(prefix=f"{SCRATCH_PREFIX}{uuid.uuid4().hex[:12]}") + try: + yield handle + finally: + # Order matters: children before parents to avoid FK conflicts. + await prisma.db.litellm_verificationtoken.delete_many( + where={ + "OR": [ + {"key_alias": {"startswith": handle.prefix}}, + {"key_name": {"startswith": handle.prefix}}, + ] + } + ) + await prisma.db.litellm_teammembership.delete_many( + where={"team_id": {"startswith": handle.prefix}} + ) + await prisma.db.litellm_organizationmembership.delete_many( + where={"user_id": {"startswith": handle.prefix}} + ) + await prisma.db.litellm_teamtable.delete_many( + where={"team_id": {"startswith": handle.prefix}} + ) + await prisma.db.litellm_usertable.delete_many( + where={"user_id": {"startswith": handle.prefix}} + ) + await prisma.db.litellm_budgettable.delete_many( + where={"budget_id": {"startswith": handle.prefix}} + ) diff --git a/tests/proxy_behavior/management/test_scratch_teardown.py b/tests/proxy_behavior/management/test_scratch_teardown.py new file mode 100644 index 00000000000..1f6bad64539 --- /dev/null +++ b/tests/proxy_behavior/management/test_scratch_teardown.py @@ -0,0 +1,41 @@ +"""Slice 5 smoke: scratch namespace fixture cleans up after itself. + +Two ordered tests. The first writes a key tagged with the scratch prefix and +asserts it lands. The second runs after the first's teardown and asserts no +scratch-namespaced rows survived. Together they prove the per-test cleanup +filter is the right shape — any leaked row will surface as a test-2 failure on +the very next run. +""" + +import pytest + +from .conftest import MASTER_KEY, SCRATCH_PREFIX + +pytestmark = pytest.mark.asyncio(loop_scope="session") + + +async def test_a_scratch_key_lands_in_db(proxy_client, prisma, scratch): + resp = await proxy_client.post( + "/key/generate", + headers={"Authorization": f"Bearer {MASTER_KEY}"}, + json={"key_alias": scratch.prefix}, + ) + assert resp.status_code == 200, resp.text + + rows = await prisma.db.litellm_verificationtoken.find_many( + where={"key_alias": scratch.prefix} + ) + assert ( + len(rows) == 1 + ), f"expected exactly one scratch-tagged token, found {len(rows)}" + + +async def test_b_scratch_namespace_is_clean(prisma): + """Runs after test_a's teardown — proves nothing leaked.""" + rows = await prisma.db.litellm_verificationtoken.find_many( + where={"key_alias": {"startswith": SCRATCH_PREFIX}} + ) + assert rows == [], ( + f"scratch teardown leaked {len(rows)} rows; first key_alias: " + f"{rows[0].key_alias if rows else None!r}" + ) diff --git a/tests/proxy_behavior/management/test_smoke.py b/tests/proxy_behavior/management/test_smoke.py index dbabd120d0a..5675a99606c 100644 --- a/tests/proxy_behavior/management/test_smoke.py +++ b/tests/proxy_behavior/management/test_smoke.py @@ -12,7 +12,7 @@ async def test_liveliness(proxy_client): assert resp.status_code == 200 -async def test_key_generate_lands_in_db(proxy_client): +async def test_key_generate_lands_in_db(proxy_client, prisma, scratch): """De-risk gate: prove the harness exercises the full stack end-to-end. A successful ``/key/generate`` requires: @@ -22,21 +22,16 @@ async def test_key_generate_lands_in_db(proxy_client): * the real ``generate_key_helper_fn`` wrote a hashed row to ``LiteLLM_VerificationToken``. - All four collapse to a single 200 + ``sk-`` token check here, with a - follow-up prisma read to prove the row landed (and that the token is the - hashed form, not the cleartext returned over the wire). + The scratch fixture tags the row with its prefix so the per-test teardown + cleans it up — keeps the proxy DB free of accumulated cruft on repeated + local runs. """ - from litellm.proxy import proxy_server from litellm.proxy.utils import hash_token - assert ( - proxy_server.prisma_client is not None - ), "FastAPI lifespan did not connect prisma — harness is wrong." - resp = await proxy_client.post( "/key/generate", headers={"Authorization": f"Bearer {MASTER_KEY}"}, - json={}, + json={"key_alias": scratch.prefix}, ) assert resp.status_code == 200, resp.text body = resp.json() @@ -44,9 +39,7 @@ async def test_key_generate_lands_in_db(proxy_client): assert cleartext_key.startswith("sk-") hashed = hash_token(cleartext_key) - row = await proxy_server.prisma_client.db.litellm_verificationtoken.find_unique( - where={"token": hashed} - ) + row = await prisma.db.litellm_verificationtoken.find_unique(where={"token": hashed}) assert row is not None, "Generated key did not land in LiteLLM_VerificationToken" assert row.token == hashed assert (