From d3cd11408204288bc4396117e24bcb099f83c73a Mon Sep 17 00:00:00 2001 From: Jack Danger Date: Sun, 18 Oct 2026 09:44:02 +0000 Subject: [PATCH] test: mark the BYOK flows async and drop re-import shadowing (ruff F811/PT017) Generated-by: termag --- .../test_r2_9_byok_flows.py | 48 +++++++------------ 1 file changed, 17 insertions(+), 31 deletions(-) diff --git a/tests/unit/proxy/management_endpoints/test_r2_9_byok_flows.py b/tests/unit/proxy/management_endpoints/test_r2_9_byok_flows.py index ea8651c216a..508dbdd1fab 100644 --- a/tests/unit/proxy/management_endpoints/test_r2_9_byok_flows.py +++ b/tests/unit/proxy/management_endpoints/test_r2_9_byok_flows.py @@ -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))