From 124b44ec2237197163fe2eedff71bd403005a92f Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Fri, 13 Mar 2026 00:02:12 -0700 Subject: [PATCH] fix(tests): update PKCE SSO tests to mock get_async_httpx_client 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 --- tests/local_testing/test_completion_cost.py | 8 +- .../proxy/management_endpoints/test_ui_sso.py | 199 +++++++++--------- 2 files changed, 98 insertions(+), 109 deletions(-) diff --git a/tests/local_testing/test_completion_cost.py b/tests/local_testing/test_completion_cost.py index 43cb236ad4..dd060a56d2 100644 --- a/tests/local_testing/test_completion_cost.py +++ b/tests/local_testing/test_completion_cost.py @@ -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) 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 ca480b8442..95ebc29fef 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py +++ b/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py @@ -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",