scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-26 06:03:19 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 35

test: mark the BYOK flows async and drop re-import shadowing (ruff F811/PT017)

Browse source 
Generated-by: termag
This commit is contained in:
Jack Danger 2026-10-18 09:44:02 +00:00
parent 45ea0fed6a
commit d3cd114082
1 changed files with 17 additions and 31 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

48
tests/unit/proxy/management_endpoints/test_r2_9_byok_flows.py
View file

@ -3,13 +3,11 @@ from unittest.mock import AsyncMock, MagicMock, patch
from litellm.types.router import GenericLiteLLMParams
async def test_team_admin_cannot_move_api_base_keeping_server_key():
import asyncio
from unittest.mock import AsyncMock, MagicMock
@pytest.mark.asyncio
async def test_team_admin_cannot_move_api_base_keeping_server_key():
from fastapi import HTTPException
from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth
from litellm.types.router import GenericLiteLLMParams
from litellm.proxy._types import LitellmUserRoles, ProxyException, UserAPIKeyAuth
from litellm.proxy.management_endpoints.model_management_endpoints import (
ModelManagementAuthChecks,
@ -20,7 +18,7 @@ async def test_team_admin_cannot_move_api_base_keeping_server_key():
model_params = MagicMock()
model_params.litellm_params = stored
with pytest.raises(Exception) as exc:
with pytest.raises(ProxyException) as exc:
await ModelManagementAuthChecks._reject_non_admin_wif_write(
model_params=model_params,
incoming_params=incoming.model_dump(exclude_none=True),
@ -36,21 +34,16 @@ async def test_team_admin_cannot_move_api_base_keeping_server_key():
@pytest.mark.asyncio
async def test_team_admin_bring_your_own_key_move_allowed():
from unittest.mock import AsyncMock
import litellm.proxy.management_endpoints.model_management_endpoints as mme
from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth
from litellm.proxy.management_endpoints.model_management_endpoints import (
ModelManagementAuthChecks,
)
from litellm.types.router import GenericLiteLLMParams
from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth
stored = GenericLiteLLMParams(api_key="server-owned-key")
incoming = GenericLiteLLMParams(
api_base="https://my-team-proxy.example", api_key="my-team-key"
)
incoming = GenericLiteLLMParams(api_base="https://my-team-proxy.example", api_key="my-team-key")
model_params = MagicMock()
model_params.litellm_params = stored
@ -70,30 +63,27 @@ async def test_team_admin_bring_your_own_key_move_allowed():
@pytest.mark.asyncio
async def test_team_admin_wif_association_still_guarded():
"""Under WIF the original server-owned-token rule keeps firing."""
from unittest.mock import AsyncMock
import litellm.proxy.management_endpoints.model_management_endpoints as mme
from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth
from litellm.proxy.management_endpoints.model_management_endpoints import (
ModelManagementAuthChecks,
)
from litellm.types.router import GenericLiteLLMParams
from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth
stored = GenericLiteLLMParams(api_key="server-key")
incoming = GenericLiteLLMParams(api_base="https://moved.example")
model_params = MagicMock()
model_params.litellm_params = stored
from litellm.types.utils import CredentialItem
with (
patch.object(
mme, "effective_server_owned_wif_fields", AsyncMock(return_value=("oidc_token_url",))
),
patch.object(mme, "effective_server_owned_wif_fields", AsyncMock(return_value=("oidc_token_url",))),
):
# Capture outside the except block: on this stack both the static-key
# harvest gate and the WIF rule can fire for the same shape; accept
# either interdiction.
guard_raised = None
try:
await ModelManagementAuthChecks._reject_non_admin_wif_write(
model_params=model_params,
@ -104,10 +94,6 @@ async def test_team_admin_wif_association_still_guarded():
prisma_client=AsyncMock(),
)
except Exception as guard_error:
# On this stack both the static-key harvest gate and the WIF
# rule can fire for the same shape; accept either interdiction.
assert ("api_key stays server-owned" in str(guard_error)) or (
"workload identity" in str(guard_error)
)
return
raise AssertionError("expected a server-owned-key interdiction to fire")
guard_raised = guard_error
assert guard_raised is not None, "expected a server-owned-key interdiction to fire"
assert ("api_key stays server-owned" in str(guard_raised)) or ("workload identity" in str(guard_raised))