refactor: Refactor oauth2_credential_exchanger to exchanger and refresher separately

PiperOrigin-RevId: 772979993

tests/unittests/auth/exchanger/test_oauth2_credential_exchanger.py

@@ -0,0 +1,220 @@
+# Copyright 2025 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import time
+from unittest.mock import Mock
+from unittest.mock import patch
+
+from authlib.oauth2.rfc6749 import OAuth2Token
+from google.adk.auth.auth_credential import AuthCredential
+from google.adk.auth.auth_credential import AuthCredentialTypes
+from google.adk.auth.auth_credential import OAuth2Auth
+from google.adk.auth.auth_schemes import OpenIdConnectWithConfig
+from google.adk.auth.exchanger.base_credential_exchanger import CredentialExchangError
+from google.adk.auth.exchanger.oauth2_credential_exchanger import OAuth2CredentialExchanger
+import pytest
+
+
+class TestOAuth2CredentialExchanger:
+  """Test suite for OAuth2CredentialExchanger."""
+
+  @pytest.mark.asyncio
+  async def test_exchange_with_existing_token(self):
+    """Test exchange method when access token already exists."""
+    scheme = OpenIdConnectWithConfig(
+        type_="openIdConnect",
+        openId_connect_url=(
+            "https://example.com/.well-known/openid_configuration"
+        ),
+        authorization_endpoint="https://example.com/auth",
+        token_endpoint="https://example.com/token",
+        scopes=["openid"],
+    )
+    credential = AuthCredential(
+        auth_type=AuthCredentialTypes.OPEN_ID_CONNECT,
+        oauth2=OAuth2Auth(
+            client_id="test_client_id",
+            client_secret="test_client_secret",
+            access_token="existing_token",
+        ),
+    )
+
+    exchanger = OAuth2CredentialExchanger()
+    result = await exchanger.exchange(credential, scheme)
+
+    # Should return the same credential since access token already exists
+    assert result == credential
+    assert result.oauth2.access_token == "existing_token"
+
+  @patch("google.adk.auth.oauth2_credential_util.OAuth2Session")
+  @pytest.mark.asyncio
+  async def test_exchange_success(self, mock_oauth2_session):
+    """Test successful token exchange."""
+    # Setup mock
+    mock_client = Mock()
+    mock_oauth2_session.return_value = mock_client
+    mock_tokens = OAuth2Token({
+        "access_token": "new_access_token",
+        "refresh_token": "new_refresh_token",
+        "expires_at": int(time.time()) + 3600,
+        "expires_in": 3600,
+    })
+    mock_client.fetch_token.return_value = mock_tokens
+
+    scheme = OpenIdConnectWithConfig(
+        type_="openIdConnect",
+        openId_connect_url=(
+            "https://example.com/.well-known/openid_configuration"
+        ),
+        authorization_endpoint="https://example.com/auth",
+        token_endpoint="https://example.com/token",
+        scopes=["openid"],
+    )
+    credential = AuthCredential(
+        auth_type=AuthCredentialTypes.OPEN_ID_CONNECT,
+        oauth2=OAuth2Auth(
+            client_id="test_client_id",
+            client_secret="test_client_secret",
+            auth_response_uri="https://example.com/callback?code=auth_code",
+            auth_code="auth_code",
+        ),
+    )
+
+    exchanger = OAuth2CredentialExchanger()
+    result = await exchanger.exchange(credential, scheme)
+
+    # Verify token exchange was successful
+    assert result.oauth2.access_token == "new_access_token"
+    assert result.oauth2.refresh_token == "new_refresh_token"
+    mock_client.fetch_token.assert_called_once()
+
+  @pytest.mark.asyncio
+  async def test_exchange_missing_auth_scheme(self):
+    """Test exchange with missing auth_scheme raises ValueError."""
+    credential = AuthCredential(
+        auth_type=AuthCredentialTypes.OPEN_ID_CONNECT,
+        oauth2=OAuth2Auth(
+            client_id="test_client_id",
+            client_secret="test_client_secret",
+        ),
+    )
+
+    exchanger = OAuth2CredentialExchanger()
+    try:
+      await exchanger.exchange(credential, None)
+      assert False, "Should have raised ValueError"
+    except CredentialExchangError as e:
+      assert "auth_scheme is required" in str(e)
+
+  @patch("google.adk.auth.oauth2_credential_util.OAuth2Session")
+  @pytest.mark.asyncio
+  async def test_exchange_no_session(self, mock_oauth2_session):
+    """Test exchange when OAuth2Session cannot be created."""
+    # Mock to return None for create_oauth2_session
+    mock_oauth2_session.return_value = None
+
+    scheme = OpenIdConnectWithConfig(
+        type_="openIdConnect",
+        openId_connect_url=(
+            "https://example.com/.well-known/openid_configuration"
+        ),
+        authorization_endpoint="https://example.com/auth",
+        token_endpoint="https://example.com/token",
+        scopes=["openid"],
+    )
+    credential = AuthCredential(
+        auth_type=AuthCredentialTypes.OPEN_ID_CONNECT,
+        oauth2=OAuth2Auth(
+            client_id="test_client_id",
+            # Missing client_secret to trigger session creation failure
+        ),
+    )
+
+    exchanger = OAuth2CredentialExchanger()
+    result = await exchanger.exchange(credential, scheme)
+
+    # Should return original credential when session creation fails
+    assert result == credential
+    assert result.oauth2.access_token is None
+
+  @patch("google.adk.auth.oauth2_credential_util.OAuth2Session")
+  @pytest.mark.asyncio
+  async def test_exchange_fetch_token_failure(self, mock_oauth2_session):
+    """Test exchange when fetch_token fails."""
+    # Setup mock to raise exception during fetch_token
+    mock_client = Mock()
+    mock_oauth2_session.return_value = mock_client
+    mock_client.fetch_token.side_effect = Exception("Token fetch failed")
+
+    scheme = OpenIdConnectWithConfig(
+        type_="openIdConnect",
+        openId_connect_url=(
+            "https://example.com/.well-known/openid_configuration"
+        ),
+        authorization_endpoint="https://example.com/auth",
+        token_endpoint="https://example.com/token",
+        scopes=["openid"],
+    )
+    credential = AuthCredential(
+        auth_type=AuthCredentialTypes.OPEN_ID_CONNECT,
+        oauth2=OAuth2Auth(
+            client_id="test_client_id",
+            client_secret="test_client_secret",
+            auth_response_uri="https://example.com/callback?code=auth_code",
+            auth_code="auth_code",
+        ),
+    )
+
+    exchanger = OAuth2CredentialExchanger()
+    result = await exchanger.exchange(credential, scheme)
+
+    # Should return original credential when fetch_token fails
+    assert result == credential
+    assert result.oauth2.access_token is None
+    mock_client.fetch_token.assert_called_once()
+
+  @pytest.mark.asyncio
+  async def test_exchange_authlib_not_available(self):
+    """Test exchange when authlib is not available."""
+    scheme = OpenIdConnectWithConfig(
+        type_="openIdConnect",
+        openId_connect_url=(
+            "https://example.com/.well-known/openid_configuration"
+        ),
+        authorization_endpoint="https://example.com/auth",
+        token_endpoint="https://example.com/token",
+        scopes=["openid"],
+    )
+    credential = AuthCredential(
+        auth_type=AuthCredentialTypes.OPEN_ID_CONNECT,
+        oauth2=OAuth2Auth(
+            client_id="test_client_id",
+            client_secret="test_client_secret",
+            auth_response_uri="https://example.com/callback?code=auth_code",
+            auth_code="auth_code",
+        ),
+    )
+
+    exchanger = OAuth2CredentialExchanger()
+
+    # Mock AUTHLIB_AVIALABLE to False
+    with patch(
+        "google.adk.auth.exchanger.oauth2_credential_exchanger.AUTHLIB_AVIALABLE",
+        False,
+    ):
+      result = await exchanger.exchange(credential, scheme)
+
+    # Should return original credential when authlib is not available
+    assert result == credential
+    assert result.oauth2.access_token is None

tests/unittests/auth/refresher/__init__.py

@@ -0,0 +1,13 @@
+# Copyright 2025 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.

tests/unittests/auth/refresher/test_oauth2_credential_refresher.py

@@ -0,0 +1,297 @@
+# Copyright 2025 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import time
+from unittest.mock import Mock
+from unittest.mock import patch
+
+from authlib.oauth2.rfc6749 import OAuth2Token
+from google.adk.auth.auth_credential import AuthCredential
+from google.adk.auth.auth_credential import AuthCredentialTypes
+from google.adk.auth.auth_credential import OAuth2Auth
+from google.adk.auth.auth_schemes import OpenIdConnectWithConfig
+from google.adk.auth.refresher.oauth2_credential_refresher import OAuth2CredentialRefresher
+import pytest
+
+
+class TestOAuth2CredentialRefresher:
+  """Test suite for OAuth2CredentialRefresher."""
+
+  @patch("google.adk.auth.refresher.oauth2_credential_refresher.OAuth2Token")
+  @pytest.mark.asyncio
+  async def test_needs_refresh_token_not_expired(self, mock_oauth2_token):
+    """Test needs_refresh when token is not expired."""
+    mock_token_instance = Mock()
+    mock_token_instance.is_expired.return_value = False
+    mock_oauth2_token.return_value = mock_token_instance
+
+    scheme = OpenIdConnectWithConfig(
+        type_="openIdConnect",
+        openId_connect_url=(
+            "https://example.com/.well-known/openid_configuration"
+        ),
+        authorization_endpoint="https://example.com/auth",
+        token_endpoint="https://example.com/token",
+        scopes=["openid"],
+    )
+    credential = AuthCredential(
+        auth_type=AuthCredentialTypes.OPEN_ID_CONNECT,
+        oauth2=OAuth2Auth(
+            client_id="test_client_id",
+            client_secret="test_client_secret",
+            access_token="existing_token",
+            expires_at=int(time.time()) + 3600,
+        ),
+    )
+
+    refresher = OAuth2CredentialRefresher()
+    needs_refresh = await refresher.is_refresh_needed(credential, scheme)
+
+    assert not needs_refresh
+
+  @patch("google.adk.auth.refresher.oauth2_credential_refresher.OAuth2Token")
+  @pytest.mark.asyncio
+  async def test_needs_refresh_token_expired(self, mock_oauth2_token):
+    """Test needs_refresh when token is expired."""
+    mock_token_instance = Mock()
+    mock_token_instance.is_expired.return_value = True
+    mock_oauth2_token.return_value = mock_token_instance
+
+    scheme = OpenIdConnectWithConfig(
+        type_="openIdConnect",
+        openId_connect_url=(
+            "https://example.com/.well-known/openid_configuration"
+        ),
+        authorization_endpoint="https://example.com/auth",
+        token_endpoint="https://example.com/token",
+        scopes=["openid"],
+    )
+    credential = AuthCredential(
+        auth_type=AuthCredentialTypes.OPEN_ID_CONNECT,
+        oauth2=OAuth2Auth(
+            client_id="test_client_id",
+            client_secret="test_client_secret",
+            access_token="existing_token",
+            expires_at=int(time.time()) - 3600,  # Expired
+        ),
+    )
+
+    refresher = OAuth2CredentialRefresher()
+    needs_refresh = await refresher.is_refresh_needed(credential, scheme)
+
+    assert needs_refresh
+
+  @patch("google.adk.auth.oauth2_credential_util.OAuth2Session")
+  @patch("google.adk.auth.oauth2_credential_util.OAuth2Token")
+  @pytest.mark.asyncio
+  async def test_refresh_token_expired_success(
+      self, mock_oauth2_token, mock_oauth2_session
+  ):
+    """Test successful token refresh when token is expired."""
+    # Setup mock token
+    mock_token_instance = Mock()
+    mock_token_instance.is_expired.return_value = True
+    mock_oauth2_token.return_value = mock_token_instance
+
+    # Setup mock session
+    mock_client = Mock()
+    mock_oauth2_session.return_value = mock_client
+    mock_tokens = OAuth2Token({
+        "access_token": "refreshed_access_token",
+        "refresh_token": "refreshed_refresh_token",
+        "expires_at": int(time.time()) + 3600,
+        "expires_in": 3600,
+    })
+    mock_client.refresh_token.return_value = mock_tokens
+
+    scheme = OpenIdConnectWithConfig(
+        type_="openIdConnect",
+        openId_connect_url=(
+            "https://example.com/.well-known/openid_configuration"
+        ),
+        authorization_endpoint="https://example.com/auth",
+        token_endpoint="https://example.com/token",
+        scopes=["openid"],
+    )
+    credential = AuthCredential(
+        auth_type=AuthCredentialTypes.OPEN_ID_CONNECT,
+        oauth2=OAuth2Auth(
+            client_id="test_client_id",
+            client_secret="test_client_secret",
+            access_token="old_token",
+            refresh_token="old_refresh_token",
+            expires_at=int(time.time()) - 3600,  # Expired
+        ),
+    )
+
+    refresher = OAuth2CredentialRefresher()
+    result = await refresher.refresh(credential, scheme)
+
+    # Verify token refresh was successful
+    assert result.oauth2.access_token == "refreshed_access_token"
+    assert result.oauth2.refresh_token == "refreshed_refresh_token"
+    mock_client.refresh_token.assert_called_once()
+
+  @pytest.mark.asyncio
+  async def test_refresh_no_oauth2_credential(self):
+    """Test refresh with no OAuth2 credential returns original."""
+    scheme = OpenIdConnectWithConfig(
+        type_="openIdConnect",
+        openId_connect_url=(
+            "https://example.com/.well-known/openid_configuration"
+        ),
+        authorization_endpoint="https://example.com/auth",
+        token_endpoint="https://example.com/token",
+        scopes=["openid"],
+    )
+    credential = AuthCredential(
+        auth_type=AuthCredentialTypes.OPEN_ID_CONNECT,
+        # No oauth2 field
+    )
+
+    refresher = OAuth2CredentialRefresher()
+    result = await refresher.refresh(credential, scheme)
+
+    assert result == credential
+
+  @pytest.mark.asyncio
+  async def test_needs_refresh_google_oauth2_json_expired(self):
+    """Test needs_refresh with Google OAuth2 JSON credential that is expired."""
+    import json
+    from unittest.mock import patch
+
+    # Mock Google OAuth2 JSON credential data
+    google_oauth2_json = json.dumps({
+        "client_id": "test_client_id",
+        "client_secret": "test_client_secret",
+        "refresh_token": "test_refresh_token",
+        "type": "authorized_user",
+    })
+
+    credential = AuthCredential(
+        auth_type=AuthCredentialTypes.OAUTH2,
+        google_oauth2_json=google_oauth2_json,
+    )
+
+    # Mock the Google Credentials class
+    with patch(
+        "google.adk.auth.refresher.oauth2_credential_refresher.Credentials"
+    ) as mock_credentials:
+      mock_google_credential = Mock()
+      mock_google_credential.expired = True
+      mock_google_credential.refresh_token = "test_refresh_token"
+      mock_credentials.from_authorized_user_info.return_value = (
+          mock_google_credential
+      )
+
+      refresher = OAuth2CredentialRefresher()
+      needs_refresh = await refresher.is_refresh_needed(credential, None)
+
+      assert needs_refresh
+
+  @pytest.mark.asyncio
+  async def test_needs_refresh_google_oauth2_json_not_expired(self):
+    """Test needs_refresh with Google OAuth2 JSON credential that is not expired."""
+    import json
+    from unittest.mock import patch
+
+    # Mock Google OAuth2 JSON credential data
+    google_oauth2_json = json.dumps({
+        "client_id": "test_client_id",
+        "client_secret": "test_client_secret",
+        "refresh_token": "test_refresh_token",
+        "type": "authorized_user",
+    })
+
+    credential = AuthCredential(
+        auth_type=AuthCredentialTypes.OAUTH2,
+        google_oauth2_json=google_oauth2_json,
+    )
+
+    # Mock the Google Credentials class
+    with patch(
+        "google.adk.auth.refresher.oauth2_credential_refresher.Credentials"
+    ) as mock_credentials:
+      mock_google_credential = Mock()
+      mock_google_credential.expired = False
+      mock_google_credential.refresh_token = "test_refresh_token"
+      mock_credentials.from_authorized_user_info.return_value = (
+          mock_google_credential
+      )
+
+      refresher = OAuth2CredentialRefresher()
+      needs_refresh = await refresher.is_refresh_needed(credential, None)
+
+      assert not needs_refresh
+
+  @pytest.mark.asyncio
+  async def test_refresh_google_oauth2_json_success(self):
+    """Test successful refresh of Google OAuth2 JSON credential."""
+    import json
+    from unittest.mock import patch
+
+    # Mock Google OAuth2 JSON credential data
+    google_oauth2_json = json.dumps({
+        "client_id": "test_client_id",
+        "client_secret": "test_client_secret",
+        "refresh_token": "test_refresh_token",
+        "type": "authorized_user",
+    })
+
+    credential = AuthCredential(
+        auth_type=AuthCredentialTypes.OAUTH2,
+        google_oauth2_json=google_oauth2_json,
+    )
+
+    # Mock the Google Credentials and Request classes
+    with patch(
+        "google.adk.auth.refresher.oauth2_credential_refresher.Credentials"
+    ) as mock_credentials:
+      with patch(
+          "google.adk.auth.refresher.oauth2_credential_refresher.Request"
+      ) as mock_request:
+        mock_google_credential = Mock()
+        mock_google_credential.expired = True
+        mock_google_credential.refresh_token = "test_refresh_token"
+        mock_google_credential.to_json.return_value = json.dumps({
+            "client_id": "test_client_id",
+            "client_secret": "test_client_secret",
+            "refresh_token": "new_refresh_token",
+            "access_token": "new_access_token",
+            "type": "authorized_user",
+        })
+        mock_credentials.from_authorized_user_info.return_value = (
+            mock_google_credential
+        )
+
+        refresher = OAuth2CredentialRefresher()
+        result = await refresher.refresh(credential, None)
+
+        mock_google_credential.refresh.assert_called_once()
+        assert (
+            result.google_oauth2_json != google_oauth2_json
+        )  # Should be updated
+
+  @pytest.mark.asyncio
+  async def test_needs_refresh_no_oauth2_credential(self):
+    """Test needs_refresh with no OAuth2 credential returns False."""
+    credential = AuthCredential(
+        auth_type=AuthCredentialTypes.HTTP,
+        # No oauth2 field
+    )
+
+    refresher = OAuth2CredentialRefresher()
+    needs_refresh = await refresher.is_refresh_needed(credential, None)
+
+    assert not needs_refresh

tests/unittests/auth/test_auth_handler.py

@@ -13,8 +13,11 @@
 # limitations under the License.
 
 import copy
+import time
+from unittest.mock import Mock
 from unittest.mock import patch
 
+from authlib.oauth2.rfc6749 import OAuth2Token
 from fastapi.openapi.models import APIKey
 from fastapi.openapi.models import APIKeyIn
 from fastapi.openapi.models import OAuth2
@@ -405,7 +408,8 @@ class TestGetAuthResponse:
 class TestParseAndStoreAuthResponse:
   """Tests for the parse_and_store_auth_response method."""
 
-  def test_non_oauth_scheme(self, auth_config_with_exchanged):
+  @pytest.mark.asyncio
+  async def test_non_oauth_scheme(self, auth_config_with_exchanged):
     """Test with a non-OAuth auth scheme."""
     # Modify the auth scheme type to be non-OAuth
     auth_config = copy.deepcopy(auth_config_with_exchanged)
@@ -416,7 +420,7 @@ class TestParseAndStoreAuthResponse:
     handler = AuthHandler(auth_config)
     state = MockState()
 
-    handler.parse_and_store_auth_response(state)
+    await handler.parse_and_store_auth_response(state)
 
     credential_key = auth_config.credential_key
     assert (
@@ -424,7 +428,10 @@ class TestParseAndStoreAuthResponse:
     )
 
   @patch("google.adk.auth.auth_handler.AuthHandler.exchange_auth_token")
-  def test_oauth_scheme(self, mock_exchange_token, auth_config_with_exchanged):
+  @pytest.mark.asyncio
+  async def test_oauth_scheme(
+      self, mock_exchange_token, auth_config_with_exchanged
+  ):
     """Test with an OAuth auth scheme."""
     mock_exchange_token.return_value = AuthCredential(
         auth_type=AuthCredentialTypes.OAUTH2,
@@ -434,7 +441,7 @@ class TestParseAndStoreAuthResponse:
     handler = AuthHandler(auth_config_with_exchanged)
     state = MockState()
 
-    handler.parse_and_store_auth_response(state)
+    await handler.parse_and_store_auth_response(state)
 
     credential_key = auth_config_with_exchanged.credential_key
     assert state["temp:" + credential_key] == mock_exchange_token.return_value
@@ -444,20 +451,20 @@ class TestParseAndStoreAuthResponse:
 class TestExchangeAuthToken:
   """Tests for the exchange_auth_token method."""
 
-  def test_token_exchange_not_supported(
+  @pytest.mark.asyncio
+  async def test_token_exchange_not_supported(
       self, auth_config_with_auth_code, monkeypatch
   ):
     """Test when token exchange is not supported."""
-    monkeypatch.setattr(
-        "google.adk.auth.oauth2_credential_fetcher.AUTHLIB_AVIALABLE", False
-    )
+    monkeypatch.setattr("google.adk.auth.auth_handler.AUTHLIB_AVIALABLE", False)
 
     handler = AuthHandler(auth_config_with_auth_code)
-    result = handler.exchange_auth_token()
+    result = await handler.exchange_auth_token()
 
     assert result == auth_config_with_auth_code.exchanged_auth_credential
 
-  def test_openid_missing_token_endpoint(
+  @pytest.mark.asyncio
+  async def test_openid_missing_token_endpoint(
       self, openid_auth_scheme, oauth2_credentials_with_auth_code
   ):
     """Test OpenID Connect without a token endpoint."""
@@ -472,11 +479,12 @@ class TestExchangeAuthToken:
     )
 
     handler = AuthHandler(config)
-    result = handler.exchange_auth_token()
+    result = await handler.exchange_auth_token()
 
     assert result == oauth2_credentials_with_auth_code
 
-  def test_oauth2_missing_token_url(
+  @pytest.mark.asyncio
+  async def test_oauth2_missing_token_url(
       self, oauth2_auth_scheme, oauth2_credentials_with_auth_code
   ):
     """Test OAuth2 without a token URL."""
@@ -491,11 +499,12 @@ class TestExchangeAuthToken:
     )
 
     handler = AuthHandler(config)
-    result = handler.exchange_auth_token()
+    result = await handler.exchange_auth_token()
 
     assert result == oauth2_credentials_with_auth_code
 
-  def test_non_oauth_scheme(self, auth_config_with_auth_code):
+  @pytest.mark.asyncio
+  async def test_non_oauth_scheme(self, auth_config_with_auth_code):
     """Test with a non-OAuth auth scheme."""
     # Modify the auth scheme type to be non-OAuth
     auth_config = copy.deepcopy(auth_config_with_auth_code)
@@ -504,11 +513,12 @@ class TestExchangeAuthToken:
     )
 
     handler = AuthHandler(auth_config)
-    result = handler.exchange_auth_token()
+    result = await handler.exchange_auth_token()
 
     assert result == auth_config.exchanged_auth_credential
 
-  def test_missing_credentials(self, oauth2_auth_scheme):
+  @pytest.mark.asyncio
+  async def test_missing_credentials(self, oauth2_auth_scheme):
     """Test with missing credentials."""
     empty_credential = AuthCredential(auth_type=AuthCredentialTypes.OAUTH2)
 
@@ -518,11 +528,12 @@ class TestExchangeAuthToken:
     )
 
     handler = AuthHandler(config)
-    result = handler.exchange_auth_token()
+    result = await handler.exchange_auth_token()
 
     assert result == empty_credential
 
-  def test_credentials_with_token(
+  @pytest.mark.asyncio
+  async def test_credentials_with_token(
       self, auth_config, oauth2_credentials_with_token
   ):
     """Test when credentials already have a token."""
@@ -533,18 +544,29 @@ class TestExchangeAuthToken:
     )
 
     handler = AuthHandler(config)
-    result = handler.exchange_auth_token()
+    result = await handler.exchange_auth_token()
 
     assert result == oauth2_credentials_with_token
 
-  @patch(
-      "google.adk.auth.oauth2_credential_util.OAuth2Session",
-      MockOAuth2Session,
-  )
-  def test_successful_token_exchange(self, auth_config_with_auth_code):
+  @patch("google.adk.auth.oauth2_credential_util.OAuth2Session")
+  @pytest.mark.asyncio
+  async def test_successful_token_exchange(
+      self, mock_oauth2_session, auth_config_with_auth_code
+  ):
     """Test a successful token exchange."""
+    # Setup mock OAuth2Session
+    mock_client = Mock()
+    mock_oauth2_session.return_value = mock_client
+    mock_tokens = OAuth2Token({
+        "access_token": "mock_access_token",
+        "refresh_token": "mock_refresh_token",
+        "expires_at": int(time.time()) + 3600,
+        "expires_in": 3600,
+    })
+    mock_client.fetch_token.return_value = mock_tokens
+
     handler = AuthHandler(auth_config_with_auth_code)
-    result = handler.exchange_auth_token()
+    result = await handler.exchange_auth_token()
 
     assert result.oauth2.access_token == "mock_access_token"
     assert result.oauth2.refresh_token == "mock_refresh_token"

tests/unittests/auth/test_oauth2_credential_fetcher.py

@@ -1,147 +0,0 @@
-# Copyright 2025 Google LLC
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-import time
-from unittest.mock import Mock
-
-from authlib.oauth2.rfc6749 import OAuth2Token
-from fastapi.openapi.models import OAuth2
-from fastapi.openapi.models import OAuthFlowAuthorizationCode
-from fastapi.openapi.models import OAuthFlows
-from google.adk.auth.auth_credential import AuthCredential
-from google.adk.auth.auth_credential import AuthCredentialTypes
-from google.adk.auth.auth_credential import OAuth2Auth
-from google.adk.auth.auth_schemes import OpenIdConnectWithConfig
-from google.adk.auth.oauth2_credential_util import create_oauth2_session
-from google.adk.auth.oauth2_credential_util import update_credential_with_tokens
-
-
-class TestOAuth2CredentialUtil:
-  """Test suite for OAuth2 credential utility functions."""
-
-  def test_create_oauth2_session_openid_connect(self):
-    """Test create_oauth2_session with OpenID Connect scheme."""
-    scheme = OpenIdConnectWithConfig(
-        type_="openIdConnect",
-        openId_connect_url=(
-            "https://example.com/.well-known/openid_configuration"
-        ),
-        authorization_endpoint="https://example.com/auth",
-        token_endpoint="https://example.com/token",
-        scopes=["openid", "profile"],
-    )
-    credential = AuthCredential(
-        auth_type=AuthCredentialTypes.OPEN_ID_CONNECT,
-        oauth2=OAuth2Auth(
-            client_id="test_client_id",
-            client_secret="test_client_secret",
-            redirect_uri="https://example.com/callback",
-            state="test_state",
-        ),
-    )
-
-    client, token_endpoint = create_oauth2_session(scheme, credential)
-
-    assert client is not None
-    assert token_endpoint == "https://example.com/token"
-    assert client.client_id == "test_client_id"
-    assert client.client_secret == "test_client_secret"
-
-  def test_create_oauth2_session_oauth2_scheme(self):
-    """Test create_oauth2_session with OAuth2 scheme."""
-    flows = OAuthFlows(
-        authorizationCode=OAuthFlowAuthorizationCode(
-            authorizationUrl="https://example.com/auth",
-            tokenUrl="https://example.com/token",
-            scopes={"read": "Read access", "write": "Write access"},
-        )
-    )
-    scheme = OAuth2(type_="oauth2", flows=flows)
-    credential = AuthCredential(
-        auth_type=AuthCredentialTypes.OAUTH2,
-        oauth2=OAuth2Auth(
-            client_id="test_client_id",
-            client_secret="test_client_secret",
-            redirect_uri="https://example.com/callback",
-        ),
-    )
-
-    client, token_endpoint = create_oauth2_session(scheme, credential)
-
-    assert client is not None
-    assert token_endpoint == "https://example.com/token"
-
-  def test_create_oauth2_session_invalid_scheme(self):
-    """Test create_oauth2_session with invalid scheme."""
-    scheme = Mock()  # Invalid scheme type
-    credential = AuthCredential(
-        auth_type=AuthCredentialTypes.OAUTH2,
-        oauth2=OAuth2Auth(
-            client_id="test_client_id",
-            client_secret="test_client_secret",
-        ),
-    )
-
-    client, token_endpoint = create_oauth2_session(scheme, credential)
-
-    assert client is None
-    assert token_endpoint is None
-
-  def test_create_oauth2_session_missing_credentials(self):
-    """Test create_oauth2_session with missing credentials."""
-    scheme = OpenIdConnectWithConfig(
-        type_="openIdConnect",
-        openId_connect_url=(
-            "https://example.com/.well-known/openid_configuration"
-        ),
-        authorization_endpoint="https://example.com/auth",
-        token_endpoint="https://example.com/token",
-        scopes=["openid"],
-    )
-    credential = AuthCredential(
-        auth_type=AuthCredentialTypes.OPEN_ID_CONNECT,
-        oauth2=OAuth2Auth(
-            client_id="test_client_id",
-            # Missing client_secret
-        ),
-    )
-
-    client, token_endpoint = create_oauth2_session(scheme, credential)
-
-    assert client is None
-    assert token_endpoint is None
-
-  def test_update_credential_with_tokens(self):
-    """Test update_credential_with_tokens function."""
-    credential = AuthCredential(
-        auth_type=AuthCredentialTypes.OPEN_ID_CONNECT,
-        oauth2=OAuth2Auth(
-            client_id="test_client_id",
-            client_secret="test_client_secret",
-        ),
-    )
-
-    tokens = OAuth2Token({
-        "access_token": "new_access_token",
-        "refresh_token": "new_refresh_token",
-        "expires_at": int(time.time()) + 3600,
-        "expires_in": 3600,
-    })
-
-    update_credential_with_tokens(credential, tokens)
-
-    assert credential.oauth2.access_token == "new_access_token"
-    assert credential.oauth2.refresh_token == "new_refresh_token"
-    assert credential.oauth2.expires_at == int(time.time()) + 3600
-    assert credential.oauth2.expires_in == 3600

tests/unittests/tools/application_integration_tool/test_integration_connector_tool.py

@@ -20,6 +20,7 @@ from google.adk.auth.auth_credential import HttpAuth
 from google.adk.auth.auth_credential import HttpCredentials
 from google.adk.tools.application_integration_tool.integration_connector_tool import IntegrationConnectorTool
 from google.adk.tools.openapi_tool.openapi_spec_parser.rest_api_tool import RestApiTool
+from google.adk.tools.openapi_tool.openapi_spec_parser.tool_auth_handler import AuthPreparationResult
 from google.genai.types import FunctionDeclaration
 from google.genai.types import Schema
 from google.genai.types import Type
@@ -50,7 +51,9 @@ def mock_rest_api_tool():
       "required": ["user_id", "page_size", "filter", "connection_name"],
   }
   mock_tool._operation_parser = mock_parser
-  mock_tool.call.return_value = {"status": "success", "data": "mock_data"}
+  mock_tool.call = mock.AsyncMock(
+      return_value={"status": "success", "data": "mock_data"}
+  )
   return mock_tool
 
 
@@ -179,9 +182,6 @@ async def test_run_with_auth_async_none_token(
       "google.adk.tools.openapi_tool.openapi_spec_parser.rest_api_tool.ToolAuthHandler.from_tool_context"
   ) as mock_from_tool_context:
     mock_tool_auth_handler_instance = mock.MagicMock()
-    mock_tool_auth_handler_instance.prepare_auth_credentials.return_value.state = (
-        "done"
-    )
     # Simulate an AuthCredential that would cause _prepare_dynamic_euc to return None
     mock_auth_credential_without_token = AuthCredential(
         auth_type=AuthCredentialTypes.HTTP,
@@ -190,8 +190,12 @@ async def test_run_with_auth_async_none_token(
             credentials=HttpCredentials(token=None),  # Token is None
         ),
     )
-    mock_tool_auth_handler_instance.prepare_auth_credentials.return_value.auth_credential = (
-        mock_auth_credential_without_token
+    mock_tool_auth_handler_instance.prepare_auth_credentials = mock.AsyncMock(
+        return_value=(
+            AuthPreparationResult(
+                state="done", auth_credential=mock_auth_credential_without_token
+            )
+        )
     )
     mock_from_tool_context.return_value = mock_tool_auth_handler_instance
 
@@ -229,18 +233,18 @@ async def test_run_with_auth_async(
       "google.adk.tools.openapi_tool.openapi_spec_parser.rest_api_tool.ToolAuthHandler.from_tool_context"
   ) as mock_from_tool_context:
     mock_tool_auth_handler_instance = mock.MagicMock()
-    mock_tool_auth_handler_instance.prepare_auth_credentials.return_value.state = (
-        "done"
-    )
-    mock_tool_auth_handler_instance.prepare_auth_credentials.return_value.state = (
-        "done"
-    )
-    mock_tool_auth_handler_instance.prepare_auth_credentials.return_value.auth_credential = AuthCredential(
-        auth_type=AuthCredentialTypes.HTTP,
-        http=HttpAuth(
-            scheme="bearer",
-            credentials=HttpCredentials(token="mocked_token"),
-        ),
+
+    mock_tool_auth_handler_instance.prepare_auth_credentials = mock.AsyncMock(
+        return_value=AuthPreparationResult(
+            state="done",
+            auth_credential=AuthCredential(
+                auth_type=AuthCredentialTypes.HTTP,
+                http=HttpAuth(
+                    scheme="bearer",
+                    credentials=HttpCredentials(token="mocked_token"),
+                ),
+            ),
+        )
     )
     mock_from_tool_context.return_value = mock_tool_auth_handler_instance
     result = await integration_tool_with_auth.run_async(

tests/unittests/tools/openapi_tool/openapi_spec_parser/test_rest_api_tool.py

@@ -14,6 +14,7 @@
 
 
 import json
+from unittest.mock import AsyncMock
 from unittest.mock import MagicMock
 from unittest.mock import patch
 
@@ -194,7 +195,8 @@ class TestRestApiTool:
   @patch(
       "google.adk.tools.openapi_tool.openapi_spec_parser.rest_api_tool.requests.request"
   )
-  def test_call_success(
+  @pytest.mark.asyncio
+  async def test_call_success(
       self,
       mock_request,
       mock_tool_context,
@@ -217,7 +219,7 @@ class TestRestApiTool:
     )
 
     # Call the method
-    result = tool.call(args={}, tool_context=mock_tool_context)
+    result = await tool.call(args={}, tool_context=mock_tool_context)
 
     # Check the result
     assert result == {"result": "success"}
@@ -225,7 +227,8 @@ class TestRestApiTool:
   @patch(
       "google.adk.tools.openapi_tool.openapi_spec_parser.rest_api_tool.requests.request"
   )
-  def test_call_auth_pending(
+  @pytest.mark.asyncio
+  async def test_call_auth_pending(
       self,
       mock_request,
       sample_endpoint,
@@ -246,12 +249,14 @@ class TestRestApiTool:
         "google.adk.tools.openapi_tool.openapi_spec_parser.rest_api_tool.ToolAuthHandler.from_tool_context"
     ) as mock_from_tool_context:
       mock_tool_auth_handler_instance = MagicMock()
-      mock_tool_auth_handler_instance.prepare_auth_credentials.return_value.state = (
-          "pending"
+      mock_prepare_result = MagicMock()
+      mock_prepare_result.state = "pending"
+      mock_tool_auth_handler_instance.prepare_auth_credentials = AsyncMock(
+          return_value=mock_prepare_result
       )
       mock_from_tool_context.return_value = mock_tool_auth_handler_instance
 
-      response = tool.call(args={}, tool_context=None)
+      response = await tool.call(args={}, tool_context=None)
       assert response == {
           "pending": True,
           "message": "Needs your authorization to access your data.",

tests/unittests/tools/openapi_tool/openapi_spec_parser/test_tool_auth_handler.py

@@ -116,7 +116,8 @@ def openid_connect_credential():
   return credential
 
 
-def test_openid_connect_no_auth_response(
+@pytest.mark.asyncio
+async def test_openid_connect_no_auth_response(
     openid_connect_scheme, openid_connect_credential
 ):
   # Setup Mock exchanger
@@ -132,12 +133,13 @@ def test_openid_connect_no_auth_response(
       credential_exchanger=mock_exchanger,
       credential_store=credential_store,
   )
-  result = handler.prepare_auth_credentials()
+  result = await handler.prepare_auth_credentials()
   assert result.state == 'pending'
   assert result.auth_credential == openid_connect_credential
 
 
-def test_openid_connect_with_auth_response(
+@pytest.mark.asyncio
+async def test_openid_connect_with_auth_response(
     openid_connect_scheme, openid_connect_credential, monkeypatch
 ):
   mock_exchanger = MockOpenIdConnectCredentialExchanger(
@@ -166,7 +168,7 @@ def test_openid_connect_with_auth_response(
       credential_exchanger=mock_exchanger,
       credential_store=credential_store,
   )
-  result = handler.prepare_auth_credentials()
+  result = await handler.prepare_auth_credentials()
   assert result.state == 'done'
   assert result.auth_credential.auth_type == AuthCredentialTypes.HTTP
   assert 'test_access_token' in result.auth_credential.http.credentials.token
@@ -178,7 +180,8 @@ def test_openid_connect_with_auth_response(
   mock_auth_handler.get_auth_response.assert_called_once()
 
 
-def test_openid_connect_existing_token(
+@pytest.mark.asyncio
+async def test_openid_connect_existing_token(
     openid_connect_scheme, openid_connect_credential
 ):
   _, existing_credential = token_to_scheme_credential(
@@ -198,16 +201,17 @@ def test_openid_connect_existing_token(
       openid_connect_credential,
       credential_store=credential_store,
   )
-  result = handler.prepare_auth_credentials()
+  result = await handler.prepare_auth_credentials()
   assert result.state == 'done'
   assert result.auth_credential == existing_credential
 
 
 @patch(
-    'google.adk.tools.openapi_tool.openapi_spec_parser.tool_auth_handler.OAuth2CredentialFetcher'
+    'google.adk.tools.openapi_tool.openapi_spec_parser.tool_auth_handler.OAuth2CredentialRefresher'
 )
-def test_openid_connect_existing_oauth2_token_refresh(
-    mock_oauth2_fetcher, openid_connect_scheme, openid_connect_credential
+@pytest.mark.asyncio
+async def test_openid_connect_existing_oauth2_token_refresh(
+    mock_oauth2_refresher, openid_connect_scheme, openid_connect_credential
 ):
   """Test that OAuth2 tokens are refreshed when existing credentials are found."""
   # Create existing OAuth2 credential
@@ -232,10 +236,13 @@ def test_openid_connect_existing_oauth2_token_refresh(
       ),
   )
 
-  # Setup mock OAuth2CredentialFetcher
-  mock_fetcher_instance = MagicMock()
-  mock_fetcher_instance.refresh.return_value = refreshed_credential
-  mock_oauth2_fetcher.return_value = mock_fetcher_instance
+  # Setup mock OAuth2CredentialRefresher
+  from unittest.mock import AsyncMock
+
+  mock_refresher_instance = MagicMock()
+  mock_refresher_instance.is_refresh_needed = AsyncMock(return_value=True)
+  mock_refresher_instance.refresh = AsyncMock(return_value=refreshed_credential)
+  mock_oauth2_refresher.return_value = mock_refresher_instance
 
   tool_context = create_mock_tool_context()
   credential_store = ToolContextCredentialStore(tool_context=tool_context)
@@ -253,13 +260,17 @@ def test_openid_connect_existing_oauth2_token_refresh(
       credential_store=credential_store,
   )
 
-  result = handler.prepare_auth_credentials()
+  result = await handler.prepare_auth_credentials()
 
-  # Verify OAuth2CredentialFetcher was called for refresh
-  mock_oauth2_fetcher.assert_called_once_with(
-      openid_connect_scheme, existing_credential
+  # Verify OAuth2CredentialRefresher was called for refresh
+  mock_oauth2_refresher.assert_called_once()
+
+  mock_refresher_instance.is_refresh_needed.assert_called_once_with(
+      existing_credential
+  )
+  mock_refresher_instance.refresh.assert_called_once_with(
+      existing_credential, openid_connect_scheme
   )
-  mock_fetcher_instance.refresh.assert_called_once()
 
   assert result.state == 'done'
   # The result should contain the refreshed credential after exchange