From e3a6666e25b4359aa5e05bf357fd5059aa2471fc Mon Sep 17 00:00:00 2001 From: Sameer Kankute Date: Thu, 25 Jun 2026 11:48:01 +0530 Subject: [PATCH] fix(auth): fix black formatting and update test mock for cli_poll_key budget lookup The get_user_object and get_team_object async calls in cli_poll_key were not mocked in the existing test, causing MagicMock await errors. Patch both functions at the auth_checks module level. Also apply black formatting to ui_sso.py which CI rejected. --- litellm/proxy/management_endpoints/ui_sso.py | 6 +++++- .../proxy/management_endpoints/test_ui_sso.py | 15 +++++++++------ 2 files changed, 14 insertions(+), 7 deletions(-) diff --git a/litellm/proxy/management_endpoints/ui_sso.py b/litellm/proxy/management_endpoints/ui_sso.py index 3b626ba6edc..56fddb70b9c 100644 --- a/litellm/proxy/management_endpoints/ui_sso.py +++ b/litellm/proxy/management_endpoints/ui_sso.py @@ -2243,7 +2243,11 @@ async def cli_poll_key( key_id: The CLI login session ID team_id: Optional team ID to assign to the JWT. If provided, must be one of user's teams. """ - from litellm.proxy.auth.auth_checks import ExperimentalUIJWTToken, get_team_object, get_user_object + from litellm.proxy.auth.auth_checks import ( + ExperimentalUIJWTToken, + get_team_object, + get_user_object, + ) from litellm.proxy.proxy_server import prisma_client, user_api_key_cache try: diff --git a/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py b/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py index 44abc7acf21..044accb341a 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py +++ b/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py @@ -2870,17 +2870,20 @@ class TestCLIKeyRegenerationFlow: with ( patch("litellm.proxy.proxy_server.user_api_key_cache", mock_cache), - patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma, + patch("litellm.proxy.proxy_server.prisma_client"), patch( "litellm.proxy.auth.auth_checks.ExperimentalUIJWTToken.get_cli_jwt_auth_token", return_value=mock_jwt_token, ) as mock_get_jwt, + patch( + "litellm.proxy.auth.auth_checks.get_user_object", + new=AsyncMock(return_value=mock_user_info), + ), + patch( + "litellm.proxy.auth.auth_checks.get_team_object", + new=AsyncMock(side_effect=Exception("no team")), + ), ): - # Mock the user lookup - mock_prisma.db.litellm_usertable.find_unique = AsyncMock( - return_value=mock_user_info - ) - # Act - Second poll with team_id result = await cli_poll_key( key_id=session_key,