tiennm99/litellm
1
0
Fork 0
mirror of https://github.com/tiennm99/litellm.git synced 2026-07-06 05:06:57 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix(tests): update PKCE SSO tests to mock get_async_httpx_client

Browse Source 
The recent commit 2a997993d4 replaced httpx.AsyncClient() with
get_async_httpx_client() in ui_sso.py, but the PKCE tests still
patched the old httpx.AsyncClient path. Updated all 10 affected
tests to mock get_async_httpx_client and removed unnecessary
context manager setup since AsyncHTTPHandler is returned directly.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
This commit is contained in:
yuneng-jiang
2026-03-13 00:02:12 -07:00
parent 5dab326d0c
commit 124b44ec22
2 changed files with 98 additions and 109 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/local_testing/test_completion_cost.py
+4 -4
View File
@@ -738,10 +738,10 @@ def test_vertex_ai_embedding_completion_cost(caplog):
text = "The quick brown fox jumps over the lazy dog."
input_tokens = litellm.token_counter(
model="vertex_ai/text-embedding-004", text=text
model="vertex_ai/textembedding-gecko", text=text
)
model_info = litellm.get_model_info(model="vertex_ai/text-embedding-004")
model_info = litellm.get_model_info(model="vertex_ai/textembedding-gecko")
print("\nExpected model info:\n{}\n\n".format(model_info))
@@ -749,7 +749,7 @@ def test_vertex_ai_embedding_completion_cost(caplog):
## CALCULATED COST
calculated_input_cost, calculated_output_cost = cost_per_token(
model="text-embedding-004",
model="textembedding-gecko",
custom_llm_provider="vertex_ai",
prompt_tokens=input_tokens,
call_type="aembedding",
@@ -824,7 +824,7 @@ async def test_completion_cost_hidden_params(sync_mode):
def test_vertex_ai_gemini_predict_cost():
model = "gemini-2.0-flash"
model = "gemini-1.5-flash"
messages = [{"role": "user", "content": "Hey, hows it going???"}]
predictive_cost = completion_cost(model=model, messages=messages)
tests/test_litellm/proxy/management_endpoints/test_ui_sso.py
+94 -105
View File
@@ -3400,7 +3400,8 @@ class TestPKCEFunctionality:
@pytest.mark.asyncio
async def test_pkce_token_exchange_basic_auth(self):
"""When include_client_id=False, client credentials go via HTTP Basic Auth."""
"""When include_client_id=False, client credentials go via HTTP Basic Auth
(encoded in the Authorization header, not via httpx.BasicAuth)."""
token_resp = {
"access_token": "tok_abc",
"id_token": None,
@@ -3417,34 +3418,22 @@ class TestPKCEFunctionality:
mock_userinfo_response.status_code = 200
mock_userinfo_response.json.return_value = userinfo_resp
captured_post_kwargs = {}
async def fake_post(*args, **kwargs):
# Verify Basic Auth is set
assert "auth" in kwargs
assert isinstance(kwargs["auth"], httpx.BasicAuth)
# Verify code_verifier is in the POST body (essential PKCE field)
post_data = kwargs.get("data", {})
assert post_data.get("code_verifier") == "verifier_abc"
# Verify redirect_uri is forwarded (required by strict OAuth providers)
assert post_data.get("redirect_uri") == "https://proxy.example.com/callback"
# Verify credentials are NOT double-sent in the POST body when using Basic Auth
assert "client_secret" not in post_data, "client_secret must not appear in POST body when using Basic Auth"
assert "client_id" not in post_data, "client_id must not appear in POST body when using Basic Auth (include_client_id=False)"
captured_post_kwargs.update(kwargs)
return mock_response
# Use separate mock clients for token exchange and userinfo —
# each httpx.AsyncClient() call gets its own independent mock.
mock_token_client = AsyncMock()
mock_token_client.__aenter__ = AsyncMock(return_value=mock_token_client)
mock_token_client.__aexit__ = AsyncMock(return_value=False)
# get_async_httpx_client returns an AsyncHTTPHandler directly (no context manager).
# _pkce_token_exchange calls it once, _get_pkce_userinfo calls it once.
mock_token_client = MagicMock()
mock_token_client.post = AsyncMock(side_effect=fake_post)
mock_userinfo_client = AsyncMock()
mock_userinfo_client.__aenter__ = AsyncMock(return_value=mock_userinfo_client)
mock_userinfo_client.__aexit__ = AsyncMock(return_value=False)
mock_userinfo_client = MagicMock()
mock_userinfo_client.get = AsyncMock(return_value=mock_userinfo_response)
with patch("litellm.proxy.management_endpoints.ui_sso.httpx.AsyncClient") as mock_client_cls:
mock_client_cls.side_effect = [mock_token_client, mock_userinfo_client]
with patch("litellm.proxy.management_endpoints.ui_sso.get_async_httpx_client") as mock_get_client:
mock_get_client.side_effect = [mock_token_client, mock_userinfo_client]
result = await SSOAuthenticationHandler._pkce_token_exchange(
authorization_code="auth_code_123",
@@ -3458,6 +3447,18 @@ class TestPKCEFunctionality:
additional_headers={},
)
# Verify Basic Auth is set in headers (not via httpx.BasicAuth kwarg)
post_headers = captured_post_kwargs.get("headers", {})
assert "Basic " in post_headers.get("Authorization", ""), "Basic Auth header must be set"
# Verify code_verifier is in the POST body (essential PKCE field)
post_data = captured_post_kwargs.get("data", {})
assert post_data.get("code_verifier") == "verifier_abc"
# Verify redirect_uri is forwarded (required by strict OAuth providers)
assert post_data.get("redirect_uri") == "https://proxy.example.com/callback"
# Verify credentials are NOT double-sent in the POST body when using Basic Auth
assert "client_secret" not in post_data, "client_secret must not appear in POST body when using Basic Auth"
assert "client_id" not in post_data, "client_id must not appear in POST body when using Basic Auth (include_client_id=False)"
assert result["access_token"] == "tok_abc"
assert result["email"] == "user@example.com"
# id_token was explicit null in token_response — the merge loop must remove it
@@ -3480,13 +3481,10 @@ class TestPKCEFunctionality:
}
userinfo_resp = {"sub": "user2", "email": "user2@example.com"}
captured_post_kwargs = {}
async def fake_post(*args, **kwargs):
assert "auth" not in kwargs, "Should NOT use Basic Auth when include_client_id=True"
data = kwargs.get("data", {})
assert "client_id" in data
assert "client_secret" in data
assert data.get("code_verifier") == "verifier_xyz", "code_verifier must be in POST body"
assert data.get("redirect_uri") == "https://proxy.example.com/callback", "redirect_uri must be forwarded"
captured_post_kwargs.update(kwargs)
mock = MagicMock()
mock.status_code = 200
mock.json.return_value = token_resp
@@ -3496,18 +3494,14 @@ class TestPKCEFunctionality:
mock_userinfo.status_code = 200
mock_userinfo.json.return_value = userinfo_resp
mock_token_client = AsyncMock()
mock_token_client.__aenter__ = AsyncMock(return_value=mock_token_client)
mock_token_client.__aexit__ = AsyncMock(return_value=False)
mock_token_client = MagicMock()
mock_token_client.post = AsyncMock(side_effect=fake_post)
mock_userinfo_client = AsyncMock()
mock_userinfo_client.__aenter__ = AsyncMock(return_value=mock_userinfo_client)
mock_userinfo_client.__aexit__ = AsyncMock(return_value=False)
mock_userinfo_client = MagicMock()
mock_userinfo_client.get = AsyncMock(return_value=mock_userinfo)
with patch("litellm.proxy.management_endpoints.ui_sso.httpx.AsyncClient") as mock_client_cls:
mock_client_cls.side_effect = [mock_token_client, mock_userinfo_client]
with patch("litellm.proxy.management_endpoints.ui_sso.get_async_httpx_client") as mock_get_client:
mock_get_client.side_effect = [mock_token_client, mock_userinfo_client]
result = await SSOAuthenticationHandler._pkce_token_exchange(
authorization_code="auth_code_456",
@@ -3521,6 +3515,15 @@ class TestPKCEFunctionality:
additional_headers={},
)
# Verify no Basic Auth header when include_client_id=True
post_headers = captured_post_kwargs.get("headers", {})
assert "Basic " not in post_headers.get("Authorization", ""), "Should NOT use Basic Auth when include_client_id=True"
data = captured_post_kwargs.get("data", {})
assert "client_id" in data
assert "client_secret" in data
assert data.get("code_verifier") == "verifier_xyz", "code_verifier must be in POST body"
assert data.get("redirect_uri") == "https://proxy.example.com/callback", "redirect_uri must be forwarded"
assert result["access_token"] == "tok_body"
assert result["sub"] == "user2"
# Verify userinfo GET used the correct Bearer token header
@@ -3536,16 +3539,14 @@ class TestPKCEFunctionality:
error_body = {"error": "invalid_grant", "error_description": "Code already used"}
with patch("litellm.proxy.management_endpoints.ui_sso.httpx.AsyncClient") as mock_client_cls:
mock_client = AsyncMock()
mock_client.__aenter__ = AsyncMock(return_value=mock_client)
mock_client.__aexit__ = AsyncMock(return_value=False)
mock_resp = MagicMock()
mock_resp.status_code = 200
mock_resp.json.return_value = error_body
mock_client.post = AsyncMock(return_value=mock_resp)
mock_client_cls.return_value = mock_client
mock_resp = MagicMock()
mock_resp.status_code = 200
mock_resp.json.return_value = error_body
mock_client = MagicMock()
mock_client.post = AsyncMock(return_value=mock_resp)
with patch("litellm.proxy.management_endpoints.ui_sso.get_async_httpx_client", return_value=mock_client):
with pytest.raises(ProxyException) as exc_info:
await SSOAuthenticationHandler._pkce_token_exchange(
authorization_code="expired_code",
@@ -3576,15 +3577,13 @@ class TestPKCEFunctionality:
).rstrip(b"=").decode()
fake_id_token = f"eyJhbGciOiJSUzI1NiJ9.{encoded_payload}.fakesig"
with patch("litellm.proxy.management_endpoints.ui_sso.httpx.AsyncClient") as mock_client_cls:
mock_client = AsyncMock()
mock_client.__aenter__ = AsyncMock(return_value=mock_client)
mock_client.__aexit__ = AsyncMock(return_value=False)
mock_fail = MagicMock()
mock_fail.status_code = 503
mock_client.get = AsyncMock(return_value=mock_fail)
mock_client_cls.return_value = mock_client
mock_fail = MagicMock()
mock_fail.status_code = 503
mock_client = MagicMock()
mock_client.get = AsyncMock(return_value=mock_fail)
with patch("litellm.proxy.management_endpoints.ui_sso.get_async_httpx_client", return_value=mock_client):
result = await SSOAuthenticationHandler._get_pkce_userinfo(
access_token="some_token",
id_token=fake_id_token,
@@ -3628,15 +3627,13 @@ class TestPKCEFunctionality:
from litellm.proxy._types import ProxyException
from litellm.proxy.management_endpoints.ui_sso import SSOAuthenticationHandler
with patch("litellm.proxy.management_endpoints.ui_sso.httpx.AsyncClient") as mock_client_cls:
mock_client = AsyncMock()
mock_client.__aenter__ = AsyncMock(return_value=mock_client)
mock_client.__aexit__ = AsyncMock(return_value=False)
mock_fail = MagicMock()
mock_fail.status_code = 503
mock_client.get = AsyncMock(return_value=mock_fail)
mock_client_cls.return_value = mock_client
mock_fail = MagicMock()
mock_fail.status_code = 503
mock_client = MagicMock()
mock_client.get = AsyncMock(return_value=mock_fail)
with patch("litellm.proxy.management_endpoints.ui_sso.get_async_httpx_client", return_value=mock_client):
with pytest.raises(ProxyException) as exc_info:
await SSOAuthenticationHandler._get_pkce_userinfo(
access_token="token",
@@ -3659,13 +3656,10 @@ class TestPKCEFunctionality:
mock_resp.status_code = 200
mock_resp.json.return_value = None # HTTP 200 with null JSON body
with patch("litellm.proxy.management_endpoints.ui_sso.httpx.AsyncClient") as mock_client_cls:
mock_client = AsyncMock()
mock_client.__aenter__ = AsyncMock(return_value=mock_client)
mock_client.__aexit__ = AsyncMock(return_value=False)
mock_client.get = AsyncMock(return_value=mock_resp)
mock_client_cls.return_value = mock_client
mock_client = MagicMock()
mock_client.get = AsyncMock(return_value=mock_resp)
with patch("litellm.proxy.management_endpoints.ui_sso.get_async_httpx_client", return_value=mock_client):
with pytest.raises(ProxyException) as exc_info:
await SSOAuthenticationHandler._get_pkce_userinfo(
access_token="access_token",
@@ -3724,12 +3718,10 @@ class TestPKCEFunctionality:
}
userinfo_resp = {"sub": "pubuser", "email": "pub@example.com"}
captured_post_kwargs = {}
async def fake_post(*args, **kwargs):
assert "auth" not in kwargs, "Public client must not use Basic Auth"
data = kwargs.get("data", {})
assert data.get("client_id") == "public_client_id"
assert "client_secret" not in data, "No secret should be sent for public client"
assert data.get("code_verifier") == "public_verifier"
captured_post_kwargs.update(kwargs)
mock = MagicMock()
mock.status_code = 200
mock.json.return_value = token_resp
@@ -3739,18 +3731,14 @@ class TestPKCEFunctionality:
mock_userinfo.status_code = 200
mock_userinfo.json.return_value = userinfo_resp
mock_token_client = AsyncMock()
mock_token_client.__aenter__ = AsyncMock(return_value=mock_token_client)
mock_token_client.__aexit__ = AsyncMock(return_value=False)
mock_token_client = MagicMock()
mock_token_client.post = AsyncMock(side_effect=fake_post)
mock_userinfo_client = AsyncMock()
mock_userinfo_client.__aenter__ = AsyncMock(return_value=mock_userinfo_client)
mock_userinfo_client.__aexit__ = AsyncMock(return_value=False)
mock_userinfo_client = MagicMock()
mock_userinfo_client.get = AsyncMock(return_value=mock_userinfo)
with patch("litellm.proxy.management_endpoints.ui_sso.httpx.AsyncClient") as mock_client_cls:
mock_client_cls.side_effect = [mock_token_client, mock_userinfo_client]
with patch("litellm.proxy.management_endpoints.ui_sso.get_async_httpx_client") as mock_get_client:
mock_get_client.side_effect = [mock_token_client, mock_userinfo_client]
result = await SSOAuthenticationHandler._pkce_token_exchange(
authorization_code="auth_pub",
@@ -3764,6 +3752,14 @@ class TestPKCEFunctionality:
additional_headers={},
)
# Verify no Basic Auth header for public client
post_headers = captured_post_kwargs.get("headers", {})
assert "Basic " not in post_headers.get("Authorization", ""), "Public client must not use Basic Auth"
data = captured_post_kwargs.get("data", {})
assert data.get("client_id") == "public_client_id"
assert "client_secret" not in data, "No secret should be sent for public client"
assert data.get("code_verifier") == "public_verifier"
assert result["access_token"] == "tok_public"
assert result["sub"] == "pubuser"
@@ -3884,13 +3880,10 @@ class TestPKCEFunctionality:
mock_response.status_code = 401
mock_response.text = "Unauthorized"
with patch("litellm.proxy.management_endpoints.ui_sso.httpx.AsyncClient") as mock_client_cls:
mock_client = AsyncMock()
mock_client.__aenter__ = AsyncMock(return_value=mock_client)
mock_client.__aexit__ = AsyncMock(return_value=False)
mock_client.post = AsyncMock(return_value=mock_response)
mock_client_cls.return_value = mock_client
mock_client = MagicMock()
mock_client.post = AsyncMock(return_value=mock_response)
with patch("litellm.proxy.management_endpoints.ui_sso.get_async_httpx_client", return_value=mock_client):
with pytest.raises(ProxyException) as exc_info:
await SSOAuthenticationHandler._pkce_token_exchange(
authorization_code="auth_code",
@@ -3993,17 +3986,15 @@ class TestPKCEFunctionality:
from litellm.proxy._types import ProxyException
from litellm.proxy.management_endpoints.ui_sso import SSOAuthenticationHandler
with patch("litellm.proxy.management_endpoints.ui_sso.httpx.AsyncClient") as mock_client_cls:
mock_client = AsyncMock()
mock_client.__aenter__ = AsyncMock(return_value=mock_client)
mock_client.__aexit__ = AsyncMock(return_value=False)
mock_resp = MagicMock()
mock_resp.status_code = 200
mock_resp.json.return_value = None # JSON null response body
mock_resp.text = "null"
mock_client.post = AsyncMock(return_value=mock_resp)
mock_client_cls.return_value = mock_client
mock_resp = MagicMock()
mock_resp.status_code = 200
mock_resp.json.return_value = None # JSON null response body
mock_resp.text = "null"
mock_client = MagicMock()
mock_client.post = AsyncMock(return_value=mock_resp)
with patch("litellm.proxy.management_endpoints.ui_sso.get_async_httpx_client", return_value=mock_client):
with pytest.raises(ProxyException) as exc_info:
await SSOAuthenticationHandler._pkce_token_exchange(
authorization_code="some_code",
@@ -4029,16 +4020,14 @@ class TestPKCEFunctionality:
body_without_token = {"token_type": "Bearer", "scope": "openid"}
with patch("litellm.proxy.management_endpoints.ui_sso.httpx.AsyncClient") as mock_client_cls:
mock_client = AsyncMock()
mock_client.__aenter__ = AsyncMock(return_value=mock_client)
mock_client.__aexit__ = AsyncMock(return_value=False)
mock_resp = MagicMock()
mock_resp.status_code = 200
mock_resp.json.return_value = body_without_token
mock_client.post = AsyncMock(return_value=mock_resp)
mock_client_cls.return_value = mock_client
mock_resp = MagicMock()
mock_resp.status_code = 200
mock_resp.json.return_value = body_without_token
mock_client = MagicMock()
mock_client.post = AsyncMock(return_value=mock_resp)
with patch("litellm.proxy.management_endpoints.ui_sso.get_async_httpx_client", return_value=mock_client):
with pytest.raises(ProxyException) as exc_info:
await SSOAuthenticationHandler._pkce_token_exchange(
authorization_code="some_code",