mirror of
https://github.com/tiennm99/litellm.git
synced 2026-08-02 06:22:48 +00:00
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:
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user