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

add modify test to perform async run

Browse source
This commit is contained in:
Harshit Jain 2026-02-13 04:04:03 +05:30
parent 847402b68d
commit e0846389e9
No known key found for this signature in database
GPG key ID: 36C392CD4415B4CF
2 changed files with 25 additions and 20 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

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

@ -11,7 +11,6 @@ Has all /sso/* routes
import asyncio
import base64
import hashlib
import json
import os
import secrets
from copy import deepcopy
@ -1723,7 +1722,7 @@ class SSOAuthenticationHandler:
if redis_usage_cache is not None:
await redis_usage_cache.async_set_cache(
key=cache_key,
value=json.dumps(code_verifier),
value=code_verifier,
ttl=600,
)
else:

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

@ -3159,10 +3159,10 @@ class TestPKCEFunctionality:
def __init__(self):
self._store = {}
def set_cache(self, key, value, **kwargs):
async def async_set_cache(self, key, value, **kwargs):
self._store[key] = value
def get_cache(self, key, **kwargs):
async def async_get_cache(self, key, **kwargs):
val = self._store.get(key)
if val is None:
return None
@ -3174,7 +3174,7 @@ class TestPKCEFunctionality:
return val
return val
def delete_cache(self, key):
async def async_delete_cache(self, key):
self._store.pop(key, None)
mock_redis = MockRedisCache()
@ -3200,7 +3200,7 @@ class TestPKCEFunctionality:
state="multi_pod_state_xyz",
generic_authorization_endpoint="https://auth.example.com/authorize",
)
mock_in_memory.set_cache.assert_not_called()
mock_in_memory.async_set_cache.assert_not_called()
# MockRedisCache is a real class; assert on state, not .assert_called_*
stored_key = "pkce_verifier:multi_pod_state_xyz"
assert stored_key in mock_redis._store
@ -3215,7 +3215,7 @@ class TestPKCEFunctionality:
)
assert "code_verifier" in token_params
assert token_params["code_verifier"] == json.loads(stored_value)
mock_in_memory.get_cache.assert_not_called()
mock_in_memory.async_get_cache.assert_not_called()
# delete_cache called; key removed (asserted below)
# Verifier consumed (single-use); key removed from "Redis"
@ -3234,19 +3234,19 @@ class TestPKCEFunctionality:
# In-memory store (simulates user_api_key_cache on one pod)
in_memory_store = {}
def set_cache(key, value, **kwargs):
async def async_set_cache(key, value, **kwargs):
in_memory_store[key] = value
def get_cache(key, **kwargs):
async def async_get_cache(key, **kwargs):
return in_memory_store.get(key)
def delete_cache(key):
async def async_delete_cache(key):
in_memory_store.pop(key, None)
mock_in_memory = MagicMock()
mock_in_memory.set_cache.side_effect = set_cache
mock_in_memory.get_cache.side_effect = get_cache
mock_in_memory.delete_cache.side_effect = delete_cache
mock_in_memory.async_set_cache = AsyncMock(side_effect=async_set_cache)
mock_in_memory.async_get_cache = AsyncMock(side_effect=async_get_cache)
mock_in_memory.async_delete_cache = AsyncMock(side_effect=async_delete_cache)
mock_sso = MagicMock()
mock_redirect_response = MagicMock()
@ -3268,9 +3268,11 @@ class TestPKCEFunctionality:
state="fallback_state_xyz",
generic_authorization_endpoint="https://auth.example.com/authorize",
)
mock_in_memory.set_cache.assert_called_once()
stored_key = mock_in_memory.set_cache.call_args.kwargs["key"]
stored_value = mock_in_memory.set_cache.call_args.kwargs["value"]
mock_in_memory.async_set_cache.assert_called_once()
stored_key = mock_in_memory.async_set_cache.call_args.kwargs["key"]
stored_value = mock_in_memory.async_set_cache.call_args.kwargs[
"value"
]
assert stored_key == "pkce_verifier:fallback_state_xyz"
assert isinstance(stored_value, str) and len(stored_value) == 43
@ -3282,8 +3284,12 @@ class TestPKCEFunctionality:
)
assert "code_verifier" in token_params
assert token_params["code_verifier"] == stored_value
mock_in_memory.get_cache.assert_called_once_with(key=stored_key)
mock_in_memory.delete_cache.assert_called_once_with(key=stored_key)
mock_in_memory.async_get_cache.assert_called_once_with(
key=stored_key
)
mock_in_memory.async_delete_cache.assert_called_once_with(
key=stored_key
)
# Verifier consumed; key removed from in-memory
assert "pkce_verifier:fallback_state_xyz" not in in_memory_store
@ -3309,8 +3315,8 @@ class TestPKCEFunctionality:
)
)
assert "code_verifier" not in token_params
mock_redis.get_cache.assert_not_called()
mock_in_memory.get_cache.assert_not_called()
mock_redis.async_get_cache.assert_not_called()
mock_in_memory.async_get_cache.assert_not_called()
# Tests for SSO user team assignment bug (Issue: SSO Users Not Added to Entra-Synced Teams on First Login)