From d2d9539aa186c806d4154b698c37e0e68b32e159 Mon Sep 17 00:00:00 2001
From: Yuneng Jiang <yuneng@berri.ai>
Date: Tue, 19 May 2026 21:24:18 -0700
Subject: [PATCH] test(proxy_behavior): connect prisma via real lifespan;
 key/generate de-risk
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Slice 3 of the management-endpoints behavior-pinning effort. The fixture now
enters the real FastAPI lifespan (proxy_startup_event) instead of just calling
initialize() — that is where prisma_client is connected, password migration is
kicked off, and the rest of the startup wiring runs.

Tests pin the loop to the session scope so the AsyncClient created in the
session fixture and the prisma connection opened in the lifespan share the
same loop as the test bodies.

New de-risk smoke: POST /key/generate with the master key returns 200, the
returned sk- token resolves to a hashed row in LiteLLM_VerificationToken, and
the cleartext token is never stored. Proves auth + handler + helper + prisma
all wire together end-to-end against a real Postgres.

Plan: https://www.notion.so/36643b8acdab8128a581ced0f6a4744d
---
 tests/proxy_behavior/management/conftest.py   | 11 ++++-
 tests/proxy_behavior/management/test_smoke.py | 48 +++++++++++++++++++
 2 files changed, 58 insertions(+), 1 deletion(-)

diff --git a/tests/proxy_behavior/management/conftest.py b/tests/proxy_behavior/management/conftest.py
index be89d6a9dc9..7deab9ba5d7 100644
--- a/tests/proxy_behavior/management/conftest.py
+++ b/tests/proxy_behavior/management/conftest.py
@@ -34,16 +34,25 @@ def _write_minimal_proxy_config() -> str:
 
 @pytest_asyncio.fixture(scope="session")
 async def proxy_app():
+    """Boot the proxy app once per session with the real FastAPI lifespan.
+
+    httpx 0.28's ASGITransport does not run the lifespan handler, so we enter
+    ``proxy_startup_event`` (the @asynccontextmanager registered as the app's
+    lifespan) directly. That handler is where ``prisma_client`` is connected
+    and the rest of the startup wiring runs.
+    """
     from litellm.proxy.proxy_server import (
         app,
         cleanup_router_config_variables,
         initialize,
+        proxy_startup_event,
     )
 
     cleanup_router_config_variables()
     config_path = _write_minimal_proxy_config()
     await initialize(config=config_path)
-    yield app
+    async with proxy_startup_event(app):
+        yield app
 
 
 @pytest_asyncio.fixture(scope="session")
diff --git a/tests/proxy_behavior/management/test_smoke.py b/tests/proxy_behavior/management/test_smoke.py
index 0481bf74f98..dbabd120d0a 100644
--- a/tests/proxy_behavior/management/test_smoke.py
+++ b/tests/proxy_behavior/management/test_smoke.py
@@ -1,6 +1,54 @@
 """Smoke tests proving the harness boots and talks to the proxy app."""
 
+import pytest
+
+from .conftest import MASTER_KEY
+
+pytestmark = pytest.mark.asyncio(loop_scope="session")
+
 
 async def test_liveliness(proxy_client):
     resp = await proxy_client.get("/health/liveliness")
     assert resp.status_code == 200
+
+
+async def test_key_generate_lands_in_db(proxy_client):
+    """De-risk gate: prove the harness exercises the full stack end-to-end.
+
+    A successful ``/key/generate`` requires:
+      * the FastAPI lifespan ran (``proxy_startup_event``),
+      * ``prisma_client`` connected,
+      * ``user_api_key_auth`` accepted the master key,
+      * 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).
+    """
+    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={},
+    )
+    assert resp.status_code == 200, resp.text
+    body = resp.json()
+    cleartext_key = body["key"]
+    assert cleartext_key.startswith("sk-")
+
+    hashed = hash_token(cleartext_key)
+    row = await proxy_server.prisma_client.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 (
+        row.token != cleartext_key
+    ), "Cleartext token stored — credential boundary broken"