Merge pull request #39917 from BerriAI/litellm_e2e_key_mgmt_route_group_coverage

test(e2e): cover key spend reset, regenerate grace period, and the llm_api_routes grant

tests/e2e/access_control/test_access_control_e2e.py

@@ -24,7 +24,7 @@ from access_control_client import (
 from e2e_config import unique_marker
 from e2e_http import Success, UnauthorizedError, UnknownApiError, unwrap
 from lifecycle import ResourceManager
-from models import ChatBody, ChatMessage, ChatResponse, LiteLLMParamsBody
+from models import ChatBody, ChatMessage, ChatResponse, EmbedBody, LiteLLMParamsBody
 from proxy_client import ProxyClient
 
 pytestmark = pytest.mark.e2e
@@ -32,6 +32,7 @@ pytestmark = pytest.mark.e2e
 ALLOWED_MODEL = "gemini-2.5-flash"
 DISALLOWED_MODEL = "gpt-5.5"
 VIRTUAL_KEY_BACKEND = "anthropic/claude-haiku-4-5-20251001"
+EMBEDDING_MODEL = "openai-text-embedding-3-small"
 
 
 class TestAccessControl:
@@ -71,6 +72,31 @@ class TestAccessControl:
             f"403 body must be a model-access denial, got: {result.body[:300]}"
         )
 
+    @pytest.mark.covers("other.auth.virtual_key.route_group_allowed")
+    def test_llm_api_routes_group_grants_every_llm_endpoint(
+        self, client: AccessControlClient, resources: ResourceManager
+    ) -> None:
+        key = client.llm_only_key()
+        resources.defer(lambda: client.delete_key(key))
+
+        chat = client.chat_status(key, ALLOWED_MODEL, f"capital of France? {unique_marker()}")
+        assert chat.status_code == 200, (
+            f"llm_api_routes key must reach /chat/completions, got {chat.status_code}: {chat.body[:300]}"
+        )
+        assert ChatResponse.model_validate_json(chat.body).choices, (
+            f"200 must carry a real completion, not an error envelope: {chat.body[:300]}"
+        )
+
+        embedding = unwrap(
+            client.proxy.embed(key, EmbedBody(model=EMBEDDING_MODEL, input=f"route group {unique_marker()}"))
+        )
+        assert embedding.model, f"llm_api_routes key reached /embeddings but got no model back: {embedding}"
+
+        denied = client.create_model_status(key, f"e2e-route-group-{unique_marker()}")
+        assert denied.status_code == 403 and ROUTE_NOT_ALLOWED_MARKER in denied.body, (
+            f"the same key must still be shut out of /model/new, got {denied.status_code}: {denied.body[:300]}"
+        )
+
     def test_llm_only_key_forbidden_from_management_route_403(
         self, client: AccessControlClient, resources: ResourceManager
     ) -> None:

tests/e2e/management/management_client.py

@@ -40,6 +40,8 @@ from models import (
     KeyListParams,
     KeyListResponse,
     KeyRegenerateBody,
+    KeyResetSpendBody,
+    KeyResetSpendResponse,
     KeyUpdateBody,
     ModelDeleteBody,
     OrgDeleteBody,
@@ -191,16 +193,26 @@ class ManagementClient:
                 response_type=NoBody,
             )
         )
-    def regenerate_key(self, key: str) -> str:
+    def regenerate_key(self, key: str, *, grace_period: str | None = None) -> str:
         return unwrap(
             self.proxy.transport.post(
                 "/key/regenerate",
                 headers=self.proxy.transport.master,
-                json=KeyRegenerateBody(key=key),
+                json=KeyRegenerateBody(key=key, grace_period=grace_period),
                 response_type=KeyGenerateResponse,
             )
         ).key
 
+    def reset_key_spend(self, key: str, reset_to: float) -> KeyResetSpendResponse:
+        return unwrap(
+            self.proxy.transport.post(
+                f"/key/{key}/reset_spend",
+                headers=self.proxy.transport.master,
+                json=KeyResetSpendBody(reset_to=reset_to),
+                response_type=KeyResetSpendResponse,
+            )
+        )
+
     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."""

tests/e2e/management/test_key_management_e2e.py

@@ -18,7 +18,7 @@ from typing import Literal
 import pytest
 
 from e2e_config import unique_marker
-from e2e_http import NoBody, unwrap
+from e2e_http import NoBody, StreamingResponse, unwrap
 from lifecycle import ResourceManager
 from management_client import ManagementClient
 from models import KeyDeleteBody, KeyGenerateBody, KeyUpdateBody
@@ -26,6 +26,9 @@ from pydantic import BaseModel
 
 pytestmark = pytest.mark.e2e
 
+TINY_BUDGET = 3e-6
+SPEND_MODEL = "claude-haiku-4-5"
+
 
 class KeyToggleBlockBody(BaseModel):
     key: str
@@ -82,6 +85,30 @@ def _generate_key(client: ManagementClient, resources: ResourceManager, body: Ke
     return key
 
 
+def _is_budget_block(outcome: StreamingResponse) -> bool:
+    return not outcome.ok and "budget_exceeded" in outcome.body
+
+
+def _spend_until_budget_blocks(client: ManagementClient, key: str) -> None:
+    for _ in range(40):
+        outcome = client.chat_status(key, SPEND_MODEL, f"spend {unique_marker()}")
+        if _is_budget_block(outcome):
+            assert outcome.status_code == 429, (
+                f"budget refusal must be 429, got {outcome.status_code}: {outcome.body[:200]}"
+            )
+            return
+        assert outcome.ok, f"paid call failed before the budget tripped ({outcome.status_code}): {outcome.body[:300]}"
+        time.sleep(2)
+    pytest.fail(f"max_budget={TINY_BUDGET} never blocked a call on the key")
+
+
+def _settled_spend(client: ManagementClient, key: str) -> float | None:
+    first = client.proxy.key_info(key).spend or 0.0
+    time.sleep(client.proxy.poll_interval)
+    second = client.proxy.key_info(key).spend or 0.0
+    return second if first > 0 and first == second else None
+
+
 def _block(client: ManagementClient, key: str) -> None:
     _ = unwrap(
         client.proxy.transport.post(
@@ -197,6 +224,32 @@ class TestKeyManagementRoutes:
             "/key/info never reported max_budget 42.0 after /key/bulk_update before the deadline",
         )
 
+    @pytest.mark.covers("other.key_mgmt.spend_reset.resets_to_value")
+    def test_reset_spend_zeroes_recorded_spend_and_lifts_the_budget_block(
+        self, client: ManagementClient, resources: ResourceManager
+    ) -> None:
+        key = _generate_key(client, resources, KeyGenerateBody(models=[SPEND_MODEL], max_budget=TINY_BUDGET))
+        _spend_until_budget_blocks(client, key)
+        recorded = _poll(
+            client, lambda: _settled_spend(client, key), "key spend never landed in /key/info before the deadline"
+        )
+
+        reset = client.reset_key_spend(key, reset_to=0.0)
+        assert reset.previous_spend == recorded, (
+            f"reset_spend reported previous_spend {reset.previous_spend}, /key/info had recorded {recorded}"
+        )
+        assert reset.spend == 0.0, f"reset_spend to 0 reported spend {reset.spend}"
+        assert client.proxy.key_info(key).spend == 0.0, "/key/info still reports spend after the reset to 0"
+
+        def call_allowed_again() -> bool | None:
+            outcome = client.chat_status(key, SPEND_MODEL, f"after reset {unique_marker()}")
+            if _is_budget_block(outcome):
+                return None
+            assert outcome.ok, f"post-reset call failed ({outcome.status_code}): {outcome.body[:300]}"
+            return True
+
+        _ = _poll(client, call_allowed_again, "the key stayed budget-blocked after its spend was reset to 0")
+
     @pytest.mark.covers("mgmt.key.generate.admin_only")
     def test_generate_forbidden_for_non_admin_key(
         self, client: ManagementClient, resources: ResourceManager

tests/e2e/management/test_management_e2e.py

@@ -12,6 +12,7 @@ from __future__ import annotations
 import math
 import time
 from collections.abc import Callable
+from typing import Final
 
 import pytest
 
@@ -42,6 +43,10 @@ from models import (
 
 pytestmark = pytest.mark.e2e
 
+REGENERATE_GRACE_PERIOD = "15s"
+REGENERATE_GRACE_SECONDS = 15.0
+
+
 def _poll[T](client: ManagementClient, attempt: Callable[[], T | None], failure: str) -> T:
     deadline = time.monotonic() + client.proxy.poll_timeout
     while time.monotonic() < deadline:
@@ -365,6 +370,36 @@ class TestKeyRegeneration:
             client, old_rejected, "old key was still accepted after regeneration (never rejected 401) at the deadline"
         )
 
+    @pytest.mark.covers("other.key_mgmt.regenerate.grace_period_honored")
+    def test_regenerate_with_grace_period_keeps_old_key_until_revoked(
+        self, client: ManagementClient, resources: ResourceManager
+    ) -> None:
+        old_key = _generate_key(client, resources, KeyGenerateBody(models=["gpt-5.5"]))
+
+        new_key = client.regenerate_key(old_key, grace_period=REGENERATE_GRACE_PERIOD)
+        resources.defer(lambda: client.proxy.delete_key(new_key))
+        revoke_at: Final = time.monotonic() + REGENERATE_GRACE_SECONDS
+        assert new_key != old_key, "regenerate returned the same key string, so no rotation happened"
+
+        def old_accepted() -> bool | None:
+            outcome = client.chat_status(old_key, "gpt-5.5", f"say hi {unique_marker()}")
+            return True if outcome.ok else None
+
+        _ = _poll(client, old_accepted, "old key was rejected 401 inside its grace period at the deadline")
+        assert time.monotonic() < revoke_at, (
+            f"old key was only accepted after its {REGENERATE_GRACE_PERIOD} grace period had elapsed"
+        )
+
+        def old_rejected() -> bool | None:
+            outcome = client.chat_status(old_key, "gpt-5.5", f"say hi {unique_marker()}")
+            return True if outcome.status_code == 401 else None
+
+        _ = _poll(
+            client,
+            old_rejected,
+            f"old key was still accepted past its {REGENERATE_GRACE_PERIOD} grace period (never 401) at the deadline",
+        )
+
 
 class TestTeamRoutes:
     @pytest.mark.covers("mgmt.team.new.persists")

tests/e2e/models.py

@@ -83,6 +83,16 @@ class KeyGenerateResponse(BaseModel):
 
 class KeyRegenerateBody(BaseModel):
     key: str
+    grace_period: str | None = None
+
+
+class KeyResetSpendBody(BaseModel):
+    reset_to: float
+
+
+class KeyResetSpendResponse(BaseModel):
+    spend: float
+    previous_spend: float
 
 
 class KeyDeleteBody(BaseModel):