test: tighten prepare_key_update_data mock and apply black

- test_prepare_key_update_data: replace bare MagicMock with
  MagicMock(spec=LiteLLM_VerificationToken) and explicitly set
  existing_key_row.metadata = {}, so reserved-field reads return real
  values instead of MagicMock-returning-MagicMock. Fixes a regression
  surfaced by the new reserved-metadata preservation logic.
- test_key_management_endpoints.py: black-format-only changes from
  recent edits.
This commit is contained in:
Ryan Crabbe
2026-04-25 09:08:55 -07:00
parent ea5c762b00
commit 7eab549190
2 changed files with 7 additions and 14 deletions
+3 -2
View File
@@ -686,12 +686,13 @@ async def test_proxy_config_update_from_db():
@pytest.mark.asyncio
async def test_prepare_key_update_data():
from litellm.proxy._types import UpdateKeyRequest
from litellm.proxy._types import LiteLLM_VerificationToken, UpdateKeyRequest
from litellm.proxy.management_endpoints.key_management_endpoints import (
prepare_key_update_data,
)
existing_key_row = MagicMock()
existing_key_row = MagicMock(spec=LiteLLM_VerificationToken)
existing_key_row.metadata = {}
data = UpdateKeyRequest(key="test_key", models=["gpt-4"], duration="120s")
updated_data = await prepare_key_update_data(data, existing_key_row)
assert "expires" in updated_data
@@ -1229,9 +1229,7 @@ async def test_update_preserves_service_account_id_when_metadata_replaced():
metadata={"service_account_id": "sa-123"},
)
result = await prepare_key_update_data(
data=data, existing_key_row=existing_key
)
result = await prepare_key_update_data(data=data, existing_key_row=existing_key)
assert result["metadata"]["service_account_id"] == "sa-123"
assert result["metadata"]["unrelated"] == "value"
@@ -1274,9 +1272,7 @@ async def test_update_allows_matching_service_account_id():
metadata={"service_account_id": "sa-123"},
)
result = await prepare_key_update_data(
data=data, existing_key_row=existing_key
)
result = await prepare_key_update_data(data=data, existing_key_row=existing_key)
assert result["metadata"]["service_account_id"] == "sa-123"
assert result["metadata"]["other"] == "value"
@@ -1334,9 +1330,7 @@ async def test_update_without_metadata_still_preserves_existing():
metadata={"service_account_id": "sa-123", "other": "kept"},
)
result = await prepare_key_update_data(
data=data, existing_key_row=existing_key
)
result = await prepare_key_update_data(data=data, existing_key_row=existing_key)
assert result["metadata"]["service_account_id"] == "sa-123"
assert result["metadata"]["other"] == "kept"
@@ -8221,9 +8215,7 @@ async def test_update_key_non_budget_rejects_cross_user_modification(monkeypatch
)
mock_prisma_client = AsyncMock()
test_hashed_token = (
"cafebabe" * 8
)
test_hashed_token = "cafebabe" * 8
mock_existing_key = MagicMock()
mock_existing_key.token = test_hashed_token