mirror of
https://github.com/BerriAI/litellm.git
synced 2026-08-28 05:25:59 +00:00
test(e2e): cover key generate and update on the Admin UI path
The two `surface: ui` cells in the coverage registry, mgmt.key.generate.happy_path and mgmt.key.update.happy_path, had no covering test. The existing key tests all call /key/generate and /key/update with the master key, which is not how the dashboard reaches those routes: an admin signs in, the proxy mints a UI session key scoped to the litellm-dashboard team, and every subsequent create or edit is written under that session key. TestDashboardKeyRoutes covers that path. The first test signs in through /v2/login, decodes the master-key-signed session JWT the way the dashboard does, and asserts the minted key carries the admin role and the dashboard team, then that it can actually read the key inventory the Virtual Keys page renders. The second edits a key under that session key and asserts both halves of the contract: /key/info reports the new models and limits with the alias untouched, and the gateway flips enforcement to match. ManagementClient grows dashboard_login plus caller-aware key_list and update_key, so a test can say who is driving a management route instead of always implying the master key. update_key returns its Result rather than raising, which lets a caller poll a route that is only transiently refusing; a freshly minted session key is briefly unauthorized while the auth cache picks up its user row.
This commit is contained in:
Yuneng Jiang
parent
f677292901
commit
d53c2c818b
3 changed files with 207 additions and 24 deletions
|
|
@ -9,8 +9,21 @@ from __future__ import annotations
|
|||
import time
|
||||
from dataclasses import dataclass
|
||||
|
||||
import jwt
|
||||
|
||||
from e2e_config import MASTER_KEY
|
||||
from proxy_client import ProxyClient
|
||||
from e2e_http import NoBody, ProbeResult, Result, StreamingResponse, Success, UnknownApiError, unwrap
|
||||
from e2e_http import (
|
||||
AuthHeaders,
|
||||
NetworkError,
|
||||
NoBody,
|
||||
ProbeResult,
|
||||
Result,
|
||||
StreamingResponse,
|
||||
Success,
|
||||
UnknownApiError,
|
||||
unwrap,
|
||||
)
|
||||
from models import (
|
||||
ChatBody,
|
||||
ChatMessage,
|
||||
|
|
@ -50,6 +63,9 @@ from models import (
|
|||
TeamNewBody,
|
||||
TeamNewResponse,
|
||||
TeamUpdateBody,
|
||||
UiLoginBody,
|
||||
UiLoginResponse,
|
||||
UiSessionClaims,
|
||||
UserDeleteBody,
|
||||
UserDeleteResponse,
|
||||
UserInfoParams,
|
||||
|
|
@ -63,38 +79,59 @@ from models import (
|
|||
|
||||
MODEL_ACCESS_DENIED_MARKER = "key_model_access_denied"
|
||||
ROUTE_NOT_ALLOWED_MARKER = "not allowed to call this route"
|
||||
DASHBOARD_SESSION_TEAM_ID = "litellm-dashboard"
|
||||
_TEAM_READY_ATTEMPTS = 15
|
||||
_TEAM_READY_SLEEP_SECONDS = 0.4
|
||||
_KEY_WRITE_ATTEMPTS = 5
|
||||
_TRANSIENT_BACKEND_MARKERS = ("connecting to redis", "name resolution")
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class DashboardSession:
|
||||
"""What a dashboard sign-in hands the Admin UI: the session key it sends as
|
||||
its bearer on every subsequent call, the claims it renders the signed-in user
|
||||
from, and where it lands the browser."""
|
||||
|
||||
session_key: str
|
||||
claims: UiSessionClaims
|
||||
redirect_url: str
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class ManagementClient:
|
||||
proxy: ProxyClient
|
||||
master_key: str
|
||||
|
||||
def llm_only_key(self) -> str:
|
||||
return self.proxy.generate_key(KeyGenerateBody(models=[], allowed_routes=["llm_api_routes"]))
|
||||
|
||||
def update_key_models(self, key: str, models: list[str]) -> None:
|
||||
last: Result[NoBody] | None = None
|
||||
for attempt in range(5):
|
||||
def update_key(self, body: KeyUpdateBody, *, caller_key: str | None = None) -> Result[NoBody]:
|
||||
"""POST /key/update. `caller_key` is who is editing: the master key by
|
||||
default, or a virtual key (the dashboard edits under the session key its
|
||||
sign-in minted, never the master key). Returns the outcome rather than
|
||||
unwrapping it, so a caller can poll a route that is only transiently
|
||||
refusing; `update_key_models` is the unwrapping shorthand."""
|
||||
headers = self.proxy.transport.master if caller_key is None else self.proxy.transport.bearer(caller_key)
|
||||
last: Result[NoBody] = NetworkError(message="/key/update was never attempted")
|
||||
for attempt in range(_KEY_WRITE_ATTEMPTS):
|
||||
last = self.proxy.transport.post(
|
||||
"/key/update",
|
||||
headers=self.proxy.transport.master,
|
||||
json=KeyUpdateBody(key=key, models=models),
|
||||
headers=headers,
|
||||
json=body,
|
||||
response_type=NoBody,
|
||||
)
|
||||
match last:
|
||||
case Success():
|
||||
return
|
||||
case UnknownApiError(body=body) if (
|
||||
"connecting to redis" in body.lower() or "name resolution" in body.lower()
|
||||
case UnknownApiError(body=error_body) if any(
|
||||
marker in error_body.lower() for marker in _TRANSIENT_BACKEND_MARKERS
|
||||
):
|
||||
time.sleep(0.5 * (attempt + 1))
|
||||
continue
|
||||
case _:
|
||||
break
|
||||
assert last is not None
|
||||
raise AssertionError(last)
|
||||
return last
|
||||
|
||||
def update_key_models(self, key: str, models: list[str]) -> None:
|
||||
_ = unwrap(self.update_key(KeyUpdateBody(key=key, models=models)))
|
||||
|
||||
def delete_key_strict(self, key: str) -> None:
|
||||
"""Strict delete for the act phase of a test: a failed delete is a hard
|
||||
|
|
@ -150,15 +187,42 @@ class ManagementClient:
|
|||
)
|
||||
).key
|
||||
|
||||
def key_list(self, key_alias: str, *, caller_key: str | None = None) -> Result[KeyListResponse]:
|
||||
"""GET /key/list, the Virtual Keys page's own inventory call. `caller_key` is
|
||||
who is asking: the master key by default, or a virtual key."""
|
||||
headers = self.proxy.transport.master if caller_key is None else self.proxy.transport.bearer(caller_key)
|
||||
return self.proxy.transport.get(
|
||||
"/key/list",
|
||||
headers=headers,
|
||||
params=KeyListParams(key_alias=key_alias),
|
||||
response_type=KeyListResponse,
|
||||
)
|
||||
|
||||
def key_alias_count(self, key_alias: str) -> int:
|
||||
return unwrap(
|
||||
self.proxy.transport.get(
|
||||
"/key/list",
|
||||
headers=self.proxy.transport.master,
|
||||
params=KeyListParams(key_alias=key_alias),
|
||||
response_type=KeyListResponse,
|
||||
return unwrap(self.key_list(key_alias)).total_count
|
||||
|
||||
def dashboard_login(self, username: str, password: str) -> DashboardSession:
|
||||
"""POST /v2/login, the call the Admin UI's sign-in form makes.
|
||||
|
||||
The proxy authenticates the credentials, mints a UI session key for the
|
||||
signed-in user, and hands it back inside a JWT signed with the master key.
|
||||
Decoding that JWT is the only way to reach the session key, and it is what
|
||||
the dashboard itself does before it can call a single management route."""
|
||||
response = unwrap(
|
||||
self.proxy.transport.post(
|
||||
"/v2/login",
|
||||
headers=AuthHeaders(),
|
||||
json=UiLoginBody(username=username, password=password),
|
||||
response_type=UiLoginResponse,
|
||||
)
|
||||
).total_count
|
||||
)
|
||||
decoded: object = jwt.decode(response.token, self.master_key, algorithms=["HS256"])
|
||||
claims = UiSessionClaims.model_validate(decoded)
|
||||
return DashboardSession(
|
||||
session_key=claims.key,
|
||||
claims=claims,
|
||||
redirect_url=response.redirect_url,
|
||||
)
|
||||
|
||||
def create_team(self, body: TeamNewBody) -> str:
|
||||
team_id = unwrap(
|
||||
|
|
@ -465,4 +529,4 @@ class ManagementClient:
|
|||
|
||||
|
||||
def build_client(proxy: ProxyClient) -> ManagementClient:
|
||||
return ManagementClient(proxy=proxy)
|
||||
return ManagementClient(proxy=proxy, master_key=MASTER_KEY)
|
||||
|
|
|
|||
|
|
@ -15,15 +15,16 @@ from collections.abc import Callable
|
|||
|
||||
import pytest
|
||||
|
||||
from e2e_config import unique_marker
|
||||
from e2e_http import StreamingResponse
|
||||
from e2e_config import UI_PASSWORD, UI_USERNAME, unique_marker
|
||||
from e2e_http import StreamingResponse, Success
|
||||
from lifecycle import ResourceManager
|
||||
from management_client import (
|
||||
DASHBOARD_SESSION_TEAM_ID,
|
||||
MODEL_ACCESS_DENIED_MARKER,
|
||||
ROUTE_NOT_ALLOWED_MARKER,
|
||||
ManagementClient,
|
||||
)
|
||||
from models import KeyGenerateBody, OrgInfoResponse, OrgNewBody, OrgUpdateBody, TagListEntry, TagNewBody, TeamNewBody, TeamUpdateBody, UserNewBody, UserUpdateBody, LiteLLMParamsBody, ModelInfoEntry
|
||||
from models import KeyGenerateBody, KeyUpdateBody, OrgInfoResponse, OrgNewBody, OrgUpdateBody, TagListEntry, TagNewBody, TeamNewBody, TeamUpdateBody, UserNewBody, UserUpdateBody, LiteLLMParamsBody, ModelInfoEntry
|
||||
|
||||
pytestmark = pytest.mark.e2e
|
||||
|
||||
|
|
@ -199,6 +200,100 @@ class TestKeyRoutes:
|
|||
return True if client.proxy.key_info(key).blocked else None
|
||||
|
||||
_ = _poll(client, blocked, "/key/info never reported the key blocked after /key/block before the deadline")
|
||||
|
||||
|
||||
class TestDashboardKeyRoutes:
|
||||
"""The /key writes as the Admin UI makes them. Signing in mints the session key
|
||||
the dashboard authenticates with, and every key an admin creates or edits in the
|
||||
browser is written under that session key rather than the master key, so these
|
||||
are the same routes the API-surface tests cover with a different caller."""
|
||||
|
||||
@pytest.mark.covers("mgmt.key.generate.happy_path")
|
||||
def test_sign_in_mints_a_session_key_that_drives_the_dashboard(
|
||||
self, client: ManagementClient, resources: ResourceManager
|
||||
) -> None:
|
||||
alias = f"e2e-mgmt-uisession-{unique_marker()}"
|
||||
_ = _generate_key(client, resources, KeyGenerateBody(models=["gemini-2.5-flash"], key_alias=alias))
|
||||
|
||||
session = client.dashboard_login(UI_USERNAME, UI_PASSWORD)
|
||||
resources.defer(lambda: client.proxy.delete_key(session.session_key))
|
||||
|
||||
assert session.claims.login_method == "username_password", (
|
||||
f"/v2/login reports login_method {session.claims.login_method!r} for a username/password sign-in"
|
||||
)
|
||||
assert session.claims.user_role == "proxy_admin", (
|
||||
f"/v2/login reports user_role {session.claims.user_role!r} for the admin credentials, expected 'proxy_admin'"
|
||||
)
|
||||
assert session.redirect_url.endswith("/ui?login=success"), (
|
||||
f"/v2/login sends the browser to {session.redirect_url!r} instead of the dashboard"
|
||||
)
|
||||
|
||||
info = client.proxy.key_info(session.session_key)
|
||||
assert info.team_id == DASHBOARD_SESSION_TEAM_ID, (
|
||||
f"the minted session key reports team_id {info.team_id!r}, expected the dashboard's "
|
||||
f"{DASHBOARD_SESSION_TEAM_ID!r}"
|
||||
)
|
||||
|
||||
def dashboard_lists_the_key() -> bool | None:
|
||||
match client.key_list(alias, caller_key=session.session_key):
|
||||
case Success(data=listing) if listing.total_count == 1:
|
||||
return True
|
||||
case _:
|
||||
return None
|
||||
|
||||
_ = _poll(
|
||||
client,
|
||||
dashboard_lists_the_key,
|
||||
f"the session key never saw {alias!r} in /key/list before the deadline, so the dashboard "
|
||||
"would render no keys",
|
||||
)
|
||||
|
||||
@pytest.mark.covers("mgmt.key.update.happy_path")
|
||||
def test_editing_a_key_from_the_dashboard_persists_and_is_enforced(
|
||||
self, client: ManagementClient, resources: ResourceManager
|
||||
) -> None:
|
||||
alias = f"e2e-mgmt-uiedit-{unique_marker()}"
|
||||
target = _generate_key(
|
||||
client,
|
||||
resources,
|
||||
KeyGenerateBody(models=["gemini-2.5-flash"], key_alias=alias, tpm_limit=100, rpm_limit=200),
|
||||
)
|
||||
_poll_chat_ok(client, target, "gemini-2.5-flash")
|
||||
_assert_model_denied(client.chat_status(target, "gpt-5.5", f"say hi {unique_marker()}"), "gpt-5.5")
|
||||
|
||||
session = client.dashboard_login(UI_USERNAME, UI_PASSWORD)
|
||||
resources.defer(lambda: client.proxy.delete_key(session.session_key))
|
||||
|
||||
def dashboard_saves_the_edit() -> bool | None:
|
||||
match client.update_key(
|
||||
KeyUpdateBody(key=target, models=["gpt-5.5"], tpm_limit=300, rpm_limit=400),
|
||||
caller_key=session.session_key,
|
||||
):
|
||||
case Success():
|
||||
return True
|
||||
case _:
|
||||
return None
|
||||
|
||||
_ = _poll(
|
||||
client,
|
||||
dashboard_saves_the_edit,
|
||||
"the dashboard session key was never accepted on /key/update before the deadline",
|
||||
)
|
||||
|
||||
info = client.proxy.key_info(target)
|
||||
assert info.models == ["gpt-5.5"], (
|
||||
f"/key/info reports models {info.models} after the dashboard edit to ['gpt-5.5']"
|
||||
)
|
||||
assert info.tpm_limit == 300, f"/key/info reports tpm_limit {info.tpm_limit} after the dashboard edit to 300"
|
||||
assert info.rpm_limit == 400, f"/key/info reports rpm_limit {info.rpm_limit} after the dashboard edit to 400"
|
||||
assert info.key_alias == alias, (
|
||||
f"the dashboard edit renamed the key to {info.key_alias!r}, it should still be {alias!r}"
|
||||
)
|
||||
|
||||
_poll_model_access_granted(client, target, "gpt-5.5")
|
||||
_poll_chat_denied(client, target, "gemini-2.5-flash")
|
||||
|
||||
|
||||
class TestKeyRegeneration:
|
||||
@pytest.mark.covers("mgmt.key.regenerate.happy_path")
|
||||
def test_regenerate_rotates_to_a_working_new_key(
|
||||
|
|
|
|||
|
|
@ -892,7 +892,10 @@ class CredentialCreateResponse(BaseModel):
|
|||
|
||||
class KeyUpdateBody(BaseModel):
|
||||
key: str
|
||||
models: list[str]
|
||||
models: list[str] | None = None
|
||||
key_alias: str | None = None
|
||||
tpm_limit: int | None = None
|
||||
rpm_limit: int | None = None
|
||||
|
||||
|
||||
class KeyBlockBody(BaseModel):
|
||||
|
|
@ -907,6 +910,27 @@ class KeyListResponse(BaseModel):
|
|||
total_count: int
|
||||
|
||||
|
||||
# ---------- admin UI session ----------
|
||||
|
||||
|
||||
class UiLoginBody(BaseModel):
|
||||
username: str
|
||||
password: str
|
||||
|
||||
|
||||
class UiLoginResponse(BaseModel):
|
||||
token: str
|
||||
redirect_url: str
|
||||
|
||||
|
||||
class UiSessionClaims(BaseModel):
|
||||
user_id: str
|
||||
key: str
|
||||
user_role: str
|
||||
login_method: Literal["sso", "username_password"]
|
||||
exp: int
|
||||
|
||||
|
||||
class TeamMemberEntry(BaseModel):
|
||||
role: Literal["admin", "user"]
|
||||
user_id: str
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue