test(proxy_behavior): connect prisma via real lifespan; key/generate de-risk

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
This commit is contained in:
Yuneng Jiang 2026-05-19 21:24:18 -07:00
parent 8af57a0873
commit d2d9539aa1
No known key found for this signature in database
2 changed files with 58 additions and 1 deletions

View file

@ -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")

View file

@ -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"