mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-26 01:12:21 +00:00
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:
parent
8af57a0873
commit
d2d9539aa1
2 changed files with 58 additions and 1 deletions
|
|
@ -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")
|
||||
|
|
|
|||
|
|
@ -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"
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue