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 <noreply@anthropic.com>

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)
 

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
+            captured_post_kwargs.update(kwargs)
-            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)"
             return mock_response
 
-        # Use separate mock clients for token exchange and userinfo —
+        # get_async_httpx_client returns an AsyncHTTPHandler directly (no context manager).
-        # each httpx.AsyncClient() call gets its own independent mock.
+        # _pkce_token_exchange calls it once, _get_pkce_userinfo calls it once.
-        mock_token_client = AsyncMock()
+        mock_token_client = MagicMock()
-        mock_token_client.__aenter__ = AsyncMock(return_value=mock_token_client)
-        mock_token_client.__aexit__ = AsyncMock(return_value=False)
         mock_token_client.post = AsyncMock(side_effect=fake_post)
 
-        mock_userinfo_client = AsyncMock()
+        mock_userinfo_client = MagicMock()
-        mock_userinfo_client.__aenter__ = AsyncMock(return_value=mock_userinfo_client)
-        mock_userinfo_client.__aexit__ = AsyncMock(return_value=False)
         mock_userinfo_client.get = AsyncMock(return_value=mock_userinfo_response)
 
-        with patch("litellm.proxy.management_endpoints.ui_sso.httpx.AsyncClient") as mock_client_cls:
+        with patch("litellm.proxy.management_endpoints.ui_sso.get_async_httpx_client") as mock_get_client:
-            mock_client_cls.side_effect = [mock_token_client, mock_userinfo_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"
+            captured_post_kwargs.update(kwargs)
-            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"
             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 = MagicMock()
-        mock_token_client.__aenter__ = AsyncMock(return_value=mock_token_client)
-        mock_token_client.__aexit__ = AsyncMock(return_value=False)
         mock_token_client.post = AsyncMock(side_effect=fake_post)
 
-        mock_userinfo_client = AsyncMock()
+        mock_userinfo_client = MagicMock()
-        mock_userinfo_client.__aenter__ = AsyncMock(return_value=mock_userinfo_client)
-        mock_userinfo_client.__aexit__ = AsyncMock(return_value=False)
         mock_userinfo_client.get = AsyncMock(return_value=mock_userinfo)
 
-        with patch("litellm.proxy.management_endpoints.ui_sso.httpx.AsyncClient") as mock_client_cls:
+        with patch("litellm.proxy.management_endpoints.ui_sso.get_async_httpx_client") as mock_get_client:
-            mock_client_cls.side_effect = [mock_token_client, mock_userinfo_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_resp = MagicMock()
-            mock_client = AsyncMock()
+        mock_resp.status_code = 200
-            mock_client.__aenter__ = AsyncMock(return_value=mock_client)
+        mock_resp.json.return_value = error_body
-            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_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_fail = MagicMock()
-            mock_client = AsyncMock()
+        mock_fail.status_code = 503
-            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_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_fail = MagicMock()
-            mock_client = AsyncMock()
+        mock_fail.status_code = 503
-            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_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 = MagicMock()
-            mock_client = AsyncMock()
+        mock_client.get = AsyncMock(return_value=mock_resp)
-            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
 
+        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"
+            captured_post_kwargs.update(kwargs)
-            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"
             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 = MagicMock()
-        mock_token_client.__aenter__ = AsyncMock(return_value=mock_token_client)
-        mock_token_client.__aexit__ = AsyncMock(return_value=False)
         mock_token_client.post = AsyncMock(side_effect=fake_post)
 
-        mock_userinfo_client = AsyncMock()
+        mock_userinfo_client = MagicMock()
-        mock_userinfo_client.__aenter__ = AsyncMock(return_value=mock_userinfo_client)
-        mock_userinfo_client.__aexit__ = AsyncMock(return_value=False)
         mock_userinfo_client.get = AsyncMock(return_value=mock_userinfo)
 
-        with patch("litellm.proxy.management_endpoints.ui_sso.httpx.AsyncClient") as mock_client_cls:
+        with patch("litellm.proxy.management_endpoints.ui_sso.get_async_httpx_client") as mock_get_client:
-            mock_client_cls.side_effect = [mock_token_client, mock_userinfo_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 = MagicMock()
-            mock_client = AsyncMock()
+        mock_client.post = AsyncMock(return_value=mock_response)
-            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
 
+        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_resp = MagicMock()
-            mock_client = AsyncMock()
+        mock_resp.status_code = 200
-            mock_client.__aenter__ = AsyncMock(return_value=mock_client)
+        mock_resp.json.return_value = None  # JSON null response body
-            mock_client.__aexit__ = AsyncMock(return_value=False)
+        mock_resp.text = "null"
-            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_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_resp = MagicMock()
-            mock_client = AsyncMock()
+        mock_resp.status_code = 200
-            mock_client.__aenter__ = AsyncMock(return_value=mock_client)
+        mock_resp.json.return_value = body_without_token
-            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_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",