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

fix(auth): fix black formatting and update test mock for cli_poll_key budget lookup

Browse source 
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.
This commit is contained in:
Sameer Kankute 2026-06-25 11:48:01 +05:30
parent 3cd6c9a2e6
commit e3a6666e25
No known key found for this signature in database
2 changed files with 14 additions and 7 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

6
litellm/proxy/management_endpoints/ui_sso.py
View file

@ -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:

15
tests/test_litellm/proxy/management_endpoints/test_ui_sso.py
View file

@ -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,