refactor: Use CallbackContext in credential service for saving/loading credential

1. credential service may be accessed by callbacks
2. plan to add load_credential and save_credential method in CallbackContext (see cl/782158513) given customer has requirement to access credential service themselves. (see https://github.com/google/adk-python/issues/1816)

It's backward compatible given CallbackContext is parent class of ToolContext

PiperOrigin-RevId: 783480378

src/google/adk/auth/credential_manager.py

@@ -16,7 +16,7 @@ from __future__ import annotations
 
 from typing import Optional
 
-from ..tools.tool_context import ToolContext
+from ..agents.callback_context import CallbackContext
 from ..utils.feature_decorator import experimental
 from .auth_credential import AuthCredential
 from .auth_credential import AuthCredentialTypes
@@ -63,7 +63,7 @@ class CredentialManager:
       )
 
       # Load and prepare credential
-      credential = await manager.load_auth_credential(tool_context)
+      credential = await manager.load_auth_credential(callback_context)
       ```
   """
 
@@ -100,11 +100,11 @@ class CredentialManager:
     """
     self._exchanger_registry.register(credential_type, exchanger_instance)
 
-  async def request_credential(self, tool_context: ToolContext) -> None:
-    tool_context.request_credential(self._auth_config)
+  async def request_credential(self, callback_context: CallbackContext) -> None:
+    callback_context.request_credential(self._auth_config)
 
   async def get_auth_credential(
-      self, tool_context: ToolContext
+      self, callback_context: CallbackContext
   ) -> Optional[AuthCredential]:
     """Load and prepare authentication credential through a structured workflow."""
 
@@ -116,14 +116,14 @@ class CredentialManager:
       return self._auth_config.raw_auth_credential
 
     # Step 3: Try to load existing processed credential
-    credential = await self._load_existing_credential(tool_context)
+    credential = await self._load_existing_credential(callback_context)
 
     # Step 4: If no existing credential, load from auth response
     # TODO instead of load from auth response, we can store auth response in
     # credential service.
     was_from_auth_response = False
     if not credential:
-      credential = await self._load_from_auth_response(tool_context)
+      credential = await self._load_from_auth_response(callback_context)
       was_from_auth_response = True
 
     # Step 5: If still no credential available, return None
@@ -134,22 +134,23 @@ class CredentialManager:
     credential, was_exchanged = await self._exchange_credential(credential)
 
     # Step 7: Refresh credential if expired
+    was_refreshed = False
     if not was_exchanged:
       credential, was_refreshed = await self._refresh_credential(credential)
 
     # Step 8: Save credential if it was modified
     if was_from_auth_response or was_exchanged or was_refreshed:
-      await self._save_credential(tool_context, credential)
+      await self._save_credential(callback_context, credential)
 
     return credential
 
   async def _load_existing_credential(
-      self, tool_context: ToolContext
+      self, callback_context: CallbackContext
   ) -> Optional[AuthCredential]:
     """Load existing credential from credential service or cached exchanged credential."""
 
     # Try loading from credential service first
-    credential = await self._load_from_credential_service(tool_context)
+    credential = await self._load_from_credential_service(callback_context)
     if credential:
       return credential
 
@@ -160,23 +161,23 @@ class CredentialManager:
     return None
 
   async def _load_from_credential_service(
-      self, tool_context: ToolContext
+      self, callback_context: CallbackContext
   ) -> Optional[AuthCredential]:
     """Load credential from credential service if available."""
-    credential_service = tool_context._invocation_context.credential_service
+    credential_service = callback_context._invocation_context.credential_service
     if credential_service:
       # Note: This should be made async in a future refactor
       # For now, assuming synchronous operation
       return await credential_service.load_credential(
-          self._auth_config, tool_context
+          self._auth_config, callback_context
       )
     return None
 
   async def _load_from_auth_response(
-      self, tool_context: ToolContext
+      self, callback_context: CallbackContext
   ) -> Optional[AuthCredential]:
-    """Load credential from auth response in tool context."""
-    return tool_context.get_auth_response(self._auth_config)
+    """Load credential from auth response in callback context."""
+    return callback_context.get_auth_response(self._auth_config)
 
   async def _exchange_credential(
       self, credential: AuthCredential
@@ -251,11 +252,13 @@ class CredentialManager:
         # Additional validation can be added here
 
   async def _save_credential(
-      self, tool_context: ToolContext, credential: AuthCredential
+      self, callback_context: CallbackContext, credential: AuthCredential
   ) -> None:
     """Save credential to credential service if available."""
-    credential_service = tool_context._invocation_context.credential_service
+    credential_service = callback_context._invocation_context.credential_service
     if credential_service:
       # Update the exchanged credential in config
       self._auth_config.exchanged_auth_credential = credential
-      await credential_service.save_credential(self._auth_config, tool_context)
+      await credential_service.save_credential(
+          self._auth_config, callback_context
+      )

src/google/adk/auth/credential_service/base_credential_service.py

@@ -18,7 +18,7 @@ from abc import ABC
 from abc import abstractmethod
 from typing import Optional
 
-from ...tools.tool_context import ToolContext
+from ...agents.callback_context import CallbackContext
 from ...utils.feature_decorator import experimental
 from ..auth_credential import AuthCredential
 from ..auth_tool import AuthConfig
@@ -33,10 +33,10 @@ class BaseCredentialService(ABC):
   async def load_credential(
       self,
       auth_config: AuthConfig,
-      tool_context: ToolContext,
+      callback_context: CallbackContext,
   ) -> Optional[AuthCredential]:
     """
-    Loads the credential by auth config and current tool context from the
+    Loads the credential by auth config and current callback context from the
     backend credential store.
 
     Args:
@@ -44,7 +44,7 @@ class BaseCredentialService(ABC):
         credential information. auth_config.get_credential_key will be used to
         build the key to load the credential.
 
-        tool_context: The context of the current invocation when the tool is
+        callback_context: The context of the current invocation when the tool is
         trying to load the credential.
 
     Returns:
@@ -56,7 +56,7 @@ class BaseCredentialService(ABC):
   async def save_credential(
       self,
       auth_config: AuthConfig,
-      tool_context: ToolContext,
+      callback_context: CallbackContext,
   ) -> None:
     """
     Saves the exchanged_auth_credential in auth config to the backend credential
@@ -67,7 +67,7 @@ class BaseCredentialService(ABC):
         credential information. auth_config.get_credential_key will be used to
         build the key to save the credential.
 
-        tool_context: The context of the current invocation when the tool is
+        callback_context: The context of the current invocation when the tool is
         trying to save the credential.
 
     Returns:

src/google/adk/auth/credential_service/in_memory_credential_service.py

@@ -18,7 +18,7 @@ from typing import Optional
 
 from typing_extensions import override
 
-from ...tools.tool_context import ToolContext
+from ...agents.callback_context import CallbackContext
 from ...utils.feature_decorator import experimental
 from ..auth_credential import AuthCredential
 from ..auth_tool import AuthConfig
@@ -37,25 +37,27 @@ class InMemoryCredentialService(BaseCredentialService):
   async def load_credential(
       self,
       auth_config: AuthConfig,
-      tool_context: ToolContext,
+      callback_context: CallbackContext,
   ) -> Optional[AuthCredential]:
-    credential_bucket = self._get_bucket_for_current_context(tool_context)
+    credential_bucket = self._get_bucket_for_current_context(callback_context)
     return credential_bucket.get(auth_config.credential_key)
 
   @override
   async def save_credential(
       self,
       auth_config: AuthConfig,
-      tool_context: ToolContext,
+      callback_context: CallbackContext,
   ) -> None:
-    credential_bucket = self._get_bucket_for_current_context(tool_context)
+    credential_bucket = self._get_bucket_for_current_context(callback_context)
     credential_bucket[auth_config.credential_key] = (
         auth_config.exchanged_auth_credential
     )
 
-  def _get_bucket_for_current_context(self, tool_context: ToolContext) -> str:
-    app_name = tool_context._invocation_context.app_name
-    user_id = tool_context._invocation_context.user_id
+  def _get_bucket_for_current_context(
+      self, callback_context: CallbackContext
+  ) -> str:
+    app_name = callback_context._invocation_context.app_name
+    user_id = callback_context._invocation_context.user_id
 
     if app_name not in self._credentials:
       self._credentials[app_name] = {}

src/google/adk/auth/credential_service/session_state_credential_service.py

@@ -18,7 +18,7 @@ from typing import Optional
 
 from typing_extensions import override
 
-from ...tools.tool_context import ToolContext
+from ...agents.callback_context import CallbackContext
 from ...utils.feature_decorator import experimental
 from ..auth_credential import AuthCredential
 from ..auth_tool import AuthConfig
@@ -36,10 +36,10 @@ class SessionStateCredentialService(BaseCredentialService):
   async def load_credential(
       self,
       auth_config: AuthConfig,
-      tool_context: ToolContext,
+      callback_context: CallbackContext,
   ) -> Optional[AuthCredential]:
     """
-    Loads the credential by auth config and current tool context from the
+    Loads the credential by auth config and current callback context from the
     backend credential store.
 
     Args:
@@ -47,20 +47,20 @@ class SessionStateCredentialService(BaseCredentialService):
         credential information. auth_config.get_credential_key will be used to
         build the key to load the credential.
 
-        tool_context: The context of the current invocation when the tool is
+        callback_context: The context of the current invocation when the tool is
         trying to load the credential.
 
     Returns:
         Optional[AuthCredential]: the credential saved in the store.
 
     """
-    return tool_context.state.get(auth_config.credential_key)
+    return callback_context.state.get(auth_config.credential_key)
 
   @override
   async def save_credential(
       self,
       auth_config: AuthConfig,
-      tool_context: ToolContext,
+      callback_context: CallbackContext,
   ) -> None:
     """
     Saves the exchanged_auth_credential in auth config to the backend credential
@@ -71,13 +71,13 @@ class SessionStateCredentialService(BaseCredentialService):
         credential information. auth_config.get_credential_key will be used to
         build the key to save the credential.
 
-        tool_context: The context of the current invocation when the tool is
+        callback_context: The context of the current invocation when the tool is
         trying to save the credential.
 
     Returns:
         None
     """
 
-    tool_context.state[auth_config.credential_key] = (
+    callback_context.state[auth_config.credential_key] = (
         auth_config.exchanged_auth_credential
     )

tests/unittests/auth/credential_service/test_in_memory_credential_service.py

@@ -17,12 +17,12 @@ from unittest.mock import Mock
 from fastapi.openapi.models import OAuth2
 from fastapi.openapi.models import OAuthFlowAuthorizationCode
 from fastapi.openapi.models import OAuthFlows
+from google.adk.agents.callback_context import CallbackContext
 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_tool import AuthConfig
 from google.adk.auth.credential_service.in_memory_credential_service import InMemoryCredentialService
-from google.adk.tools.tool_context import ToolContext
 import pytest
 
 
@@ -69,9 +69,9 @@ class TestInMemoryCredentialService:
     )
 
   @pytest.fixture
-  def tool_context(self):
-    """Create a mock ToolContext for testing."""
-    mock_context = Mock(spec=ToolContext)
+  def callback_context(self):
+    """Create a mock CallbackContext for testing."""
+    mock_context = Mock(spec=CallbackContext)
     mock_invocation_context = Mock()
     mock_invocation_context.app_name = "test_app"
     mock_invocation_context.user_id = "test_user"
@@ -79,9 +79,9 @@ class TestInMemoryCredentialService:
     return mock_context
 
   @pytest.fixture
-  def another_tool_context(self):
-    """Create another mock ToolContext with different app/user for testing isolation."""
-    mock_context = Mock(spec=ToolContext)
+  def another_callback_context(self):
+    """Create another mock CallbackContext with different app/user for testing isolation."""
+    mock_context = Mock(spec=CallbackContext)
     mock_invocation_context = Mock()
     mock_invocation_context.app_name = "another_app"
     mock_invocation_context.user_id = "another_user"
@@ -95,22 +95,26 @@ class TestInMemoryCredentialService:
 
   @pytest.mark.asyncio
   async def test_load_credential_not_found(
-      self, credential_service, auth_config, tool_context
+      self, credential_service, auth_config, callback_context
   ):
     """Test loading a credential that doesn't exist returns None."""
-    result = await credential_service.load_credential(auth_config, tool_context)
+    result = await credential_service.load_credential(
+        auth_config, callback_context
+    )
     assert result is None
 
   @pytest.mark.asyncio
   async def test_save_and_load_credential(
-      self, credential_service, auth_config, tool_context
+      self, credential_service, auth_config, callback_context
   ):
     """Test saving and then loading a credential."""
     # Save the credential
-    await credential_service.save_credential(auth_config, tool_context)
+    await credential_service.save_credential(auth_config, callback_context)
 
     # Load the credential
-    result = await credential_service.load_credential(auth_config, tool_context)
+    result = await credential_service.load_credential(
+        auth_config, callback_context
+    )
 
     # Verify the credential was saved and loaded correctly
     assert result is not None
@@ -120,11 +124,15 @@ class TestInMemoryCredentialService:
 
   @pytest.mark.asyncio
   async def test_save_credential_updates_existing(
-      self, credential_service, auth_config, tool_context, oauth2_credentials
+      self,
+      credential_service,
+      auth_config,
+      callback_context,
+      oauth2_credentials,
   ):
     """Test that saving a credential updates an existing one."""
     # Save initial credential
-    await credential_service.save_credential(auth_config, tool_context)
+    await credential_service.save_credential(auth_config, callback_context)
 
     # Create a new credential and update the auth_config
     new_credential = AuthCredential(
@@ -138,35 +146,43 @@ class TestInMemoryCredentialService:
     auth_config.exchanged_auth_credential = new_credential
 
     # Save the updated credential
-    await credential_service.save_credential(auth_config, tool_context)
+    await credential_service.save_credential(auth_config, callback_context)
 
     # Load and verify the credential was updated
-    result = await credential_service.load_credential(auth_config, tool_context)
+    result = await credential_service.load_credential(
+        auth_config, callback_context
+    )
     assert result is not None
     assert result.oauth2.client_id == "updated_client_id"
     assert result.oauth2.client_secret == "updated_client_secret"
 
   @pytest.mark.asyncio
   async def test_credentials_isolated_by_context(
-      self, credential_service, auth_config, tool_context, another_tool_context
+      self,
+      credential_service,
+      auth_config,
+      callback_context,
+      another_callback_context,
   ):
     """Test that credentials are isolated between different app/user contexts."""
     # Save credential in first context
-    await credential_service.save_credential(auth_config, tool_context)
+    await credential_service.save_credential(auth_config, callback_context)
 
     # Try to load from another context
     result = await credential_service.load_credential(
-        auth_config, another_tool_context
+        auth_config, another_callback_context
     )
     assert result is None
 
     # Verify original context still has the credential
-    result = await credential_service.load_credential(auth_config, tool_context)
+    result = await credential_service.load_credential(
+        auth_config, callback_context
+    )
     assert result is not None
 
   @pytest.mark.asyncio
   async def test_multiple_credentials_same_context(
-      self, credential_service, tool_context, oauth2_auth_scheme
+      self, credential_service, callback_context, oauth2_auth_scheme
   ):
     """Test storing multiple credentials in the same context with different keys."""
     # Create two different auth configs with different credential keys
@@ -203,15 +219,15 @@ class TestInMemoryCredentialService:
     )
 
     # Save both credentials
-    await credential_service.save_credential(auth_config1, tool_context)
-    await credential_service.save_credential(auth_config2, tool_context)
+    await credential_service.save_credential(auth_config1, callback_context)
+    await credential_service.save_credential(auth_config2, callback_context)
 
     # Load and verify both credentials
     result1 = await credential_service.load_credential(
-        auth_config1, tool_context
+        auth_config1, callback_context
     )
     result2 = await credential_service.load_credential(
-        auth_config2, tool_context
+        auth_config2, callback_context
     )
 
     assert result1 is not None
@@ -220,10 +236,12 @@ class TestInMemoryCredentialService:
     assert result2.oauth2.client_id == "client2"
 
   def test_get_bucket_for_current_context_creates_nested_structure(
-      self, credential_service, tool_context
+      self, credential_service, callback_context
   ):
     """Test that _get_bucket_for_current_context creates the proper nested structure."""
-    storage = credential_service._get_bucket_for_current_context(tool_context)
+    storage = credential_service._get_bucket_for_current_context(
+        callback_context
+    )
 
     # Verify the nested structure was created
     assert "test_app" in credential_service._credentials
@@ -232,27 +250,33 @@ class TestInMemoryCredentialService:
     assert storage is credential_service._credentials["test_app"]["test_user"]
 
   def test_get_bucket_for_current_context_reuses_existing(
-      self, credential_service, tool_context
+      self, credential_service, callback_context
   ):
     """Test that _get_bucket_for_current_context reuses existing structure."""
     # Create initial structure
-    storage1 = credential_service._get_bucket_for_current_context(tool_context)
+    storage1 = credential_service._get_bucket_for_current_context(
+        callback_context
+    )
     storage1["test_key"] = "test_value"
 
     # Get storage again
-    storage2 = credential_service._get_bucket_for_current_context(tool_context)
+    storage2 = credential_service._get_bucket_for_current_context(
+        callback_context
+    )
 
     # Verify it's the same storage instance
     assert storage1 is storage2
     assert storage2["test_key"] == "test_value"
 
   def test_get_storage_different_apps(
-      self, credential_service, tool_context, another_tool_context
+      self, credential_service, callback_context, another_callback_context
   ):
     """Test that different apps get different storage instances."""
-    storage1 = credential_service._get_bucket_for_current_context(tool_context)
+    storage1 = credential_service._get_bucket_for_current_context(
+        callback_context
+    )
     storage2 = credential_service._get_bucket_for_current_context(
-        another_tool_context
+        another_callback_context
     )
 
     # Verify they are different storage instances
@@ -270,26 +294,26 @@ class TestInMemoryCredentialService:
   ):
     """Test that the same user in different apps get isolated storage."""
     # Create two contexts with same user but different apps
-    context1 = Mock(spec=ToolContext)
+    context1 = Mock(spec=CallbackContext)
     mock_invocation_context1 = Mock()
     mock_invocation_context1.app_name = "app1"
     mock_invocation_context1.user_id = "same_user"
     context1._invocation_context = mock_invocation_context1
 
-    context2 = Mock(spec=ToolContext)
+    context2 = Mock(spec=CallbackContext)
     mock_invocation_context2 = Mock()
     mock_invocation_context2.app_name = "app2"
     mock_invocation_context2.user_id = "same_user"
     context2._invocation_context = mock_invocation_context2
 
-    # Save credential in app1
+    # Save credential in first app
     await credential_service.save_credential(auth_config, context1)
 
-    # Try to load from app2 (should not find it)
+    # Try to load from second app
     result = await credential_service.load_credential(auth_config, context2)
     assert result is None
 
-    # Verify app1 still has the credential
+    # Verify first app still has the credential
     result = await credential_service.load_credential(auth_config, context1)
     assert result is not None
 
@@ -299,25 +323,25 @@ class TestInMemoryCredentialService:
   ):
     """Test that different users in the same app get isolated storage."""
     # Create two contexts with same app but different users
-    context1 = Mock(spec=ToolContext)
+    context1 = Mock(spec=CallbackContext)
     mock_invocation_context1 = Mock()
     mock_invocation_context1.app_name = "same_app"
     mock_invocation_context1.user_id = "user1"
     context1._invocation_context = mock_invocation_context1
 
-    context2 = Mock(spec=ToolContext)
+    context2 = Mock(spec=CallbackContext)
     mock_invocation_context2 = Mock()
     mock_invocation_context2.app_name = "same_app"
     mock_invocation_context2.user_id = "user2"
     context2._invocation_context = mock_invocation_context2
 
-    # Save credential for user1
+    # Save credential for first user
     await credential_service.save_credential(auth_config, context1)
 
-    # Try to load for user2 (should not find it)
+    # Try to load for second user
     result = await credential_service.load_credential(auth_config, context2)
     assert result is None
 
-    # Verify user1 still has the credential
+    # Verify first user still has the credential
     result = await credential_service.load_credential(auth_config, context1)
     assert result is not None

tests/unittests/auth/credential_service/test_session_state_credential_service.py

@@ -17,12 +17,12 @@ from unittest.mock import Mock
 from fastapi.openapi.models import OAuth2
 from fastapi.openapi.models import OAuthFlowAuthorizationCode
 from fastapi.openapi.models import OAuthFlows
+from google.adk.agents.callback_context import CallbackContext
 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_tool import AuthConfig
 from google.adk.auth.credential_service.session_state_credential_service import SessionStateCredentialService
-from google.adk.tools.tool_context import ToolContext
 import pytest
 
 
@@ -69,39 +69,43 @@ class TestSessionStateCredentialService:
     )
 
   @pytest.fixture
-  def tool_context(self):
-    """Create a mock ToolContext for testing."""
-    mock_context = Mock(spec=ToolContext)
+  def callback_context(self):
+    """Create a mock CallbackContext for testing."""
+    mock_context = Mock(spec=CallbackContext)
     # Create a state dictionary that behaves like session state
     mock_context.state = {}
     return mock_context
 
   @pytest.fixture
-  def another_tool_context(self):
-    """Create another mock ToolContext with different state for testing isolation."""
-    mock_context = Mock(spec=ToolContext)
+  def another_callback_context(self):
+    """Create another mock CallbackContext with different state for testing isolation."""
+    mock_context = Mock(spec=CallbackContext)
     # Create a separate state dictionary to simulate different session
     mock_context.state = {}
     return mock_context
 
   @pytest.mark.asyncio
   async def test_load_credential_not_found(
-      self, credential_service, auth_config, tool_context
+      self, credential_service, auth_config, callback_context
   ):
     """Test loading a credential that doesn't exist returns None."""
-    result = await credential_service.load_credential(auth_config, tool_context)
+    result = await credential_service.load_credential(
+        auth_config, callback_context
+    )
     assert result is None
 
   @pytest.mark.asyncio
   async def test_save_and_load_credential(
-      self, credential_service, auth_config, tool_context
+      self, credential_service, auth_config, callback_context
   ):
     """Test saving and then loading a credential."""
     # Save the credential
-    await credential_service.save_credential(auth_config, tool_context)
+    await credential_service.save_credential(auth_config, callback_context)
 
     # Load the credential
-    result = await credential_service.load_credential(auth_config, tool_context)
+    result = await credential_service.load_credential(
+        auth_config, callback_context
+    )
 
     # Verify the credential was saved and loaded correctly
     assert result is not None
@@ -111,11 +115,15 @@ class TestSessionStateCredentialService:
 
   @pytest.mark.asyncio
   async def test_save_credential_updates_existing(
-      self, credential_service, auth_config, tool_context, oauth2_credentials
+      self,
+      credential_service,
+      auth_config,
+      callback_context,
+      oauth2_credentials,
   ):
     """Test that saving a credential updates an existing one."""
     # Save initial credential
-    await credential_service.save_credential(auth_config, tool_context)
+    await credential_service.save_credential(auth_config, callback_context)
 
     # Create a new credential and update the auth_config
     new_credential = AuthCredential(
@@ -129,35 +137,43 @@ class TestSessionStateCredentialService:
     auth_config.exchanged_auth_credential = new_credential
 
     # Save the updated credential
-    await credential_service.save_credential(auth_config, tool_context)
+    await credential_service.save_credential(auth_config, callback_context)
 
     # Load and verify the credential was updated
-    result = await credential_service.load_credential(auth_config, tool_context)
+    result = await credential_service.load_credential(
+        auth_config, callback_context
+    )
     assert result is not None
     assert result.oauth2.client_id == "updated_client_id"
     assert result.oauth2.client_secret == "updated_client_secret"
 
   @pytest.mark.asyncio
   async def test_credentials_isolated_by_context(
-      self, credential_service, auth_config, tool_context, another_tool_context
+      self,
+      credential_service,
+      auth_config,
+      callback_context,
+      another_callback_context,
   ):
-    """Test that credentials are isolated between different tool contexts."""
+    """Test that credentials are isolated between different callback contexts."""
     # Save credential in first context
-    await credential_service.save_credential(auth_config, tool_context)
+    await credential_service.save_credential(auth_config, callback_context)
 
     # Try to load from another context (should not find it)
     result = await credential_service.load_credential(
-        auth_config, another_tool_context
+        auth_config, another_callback_context
     )
     assert result is None
 
     # Verify original context still has the credential
-    result = await credential_service.load_credential(auth_config, tool_context)
+    result = await credential_service.load_credential(
+        auth_config, callback_context
+    )
     assert result is not None
 
   @pytest.mark.asyncio
   async def test_multiple_credentials_same_context(
-      self, credential_service, tool_context, oauth2_auth_scheme
+      self, credential_service, callback_context, oauth2_auth_scheme
   ):
     """Test storing multiple credentials in the same context with different keys."""
     # Create two different auth configs with different credential keys
@@ -194,15 +210,15 @@ class TestSessionStateCredentialService:
     )
 
     # Save both credentials
-    await credential_service.save_credential(auth_config1, tool_context)
-    await credential_service.save_credential(auth_config2, tool_context)
+    await credential_service.save_credential(auth_config1, callback_context)
+    await credential_service.save_credential(auth_config2, callback_context)
 
     # Load and verify both credentials
     result1 = await credential_service.load_credential(
-        auth_config1, tool_context
+        auth_config1, callback_context
     )
     result2 = await credential_service.load_credential(
-        auth_config2, tool_context
+        auth_config2, callback_context
     )
 
     assert result1 is not None
@@ -212,144 +228,161 @@ class TestSessionStateCredentialService:
 
   @pytest.mark.asyncio
   async def test_save_credential_with_none_exchanged_credential(
-      self, credential_service, auth_config, tool_context
+      self, credential_service, auth_config, callback_context
   ):
-    """Test saving when exchanged_auth_credential is None."""
-    # Set exchanged credential to None
+    """Test that saving a credential with None exchanged_auth_credential stores None."""
+    # Set exchanged_auth_credential to None
     auth_config.exchanged_auth_credential = None
 
-    # Save the credential (should save None)
-    await credential_service.save_credential(auth_config, tool_context)
+    # Save the credential
+    await credential_service.save_credential(auth_config, callback_context)
 
-    # Load and verify None was saved
-    result = await credential_service.load_credential(auth_config, tool_context)
+    # Load and verify None was stored
+    result = await credential_service.load_credential(
+        auth_config, callback_context
+    )
     assert result is None
 
   @pytest.mark.asyncio
   async def test_load_credential_with_empty_credential_key(
-      self, credential_service, auth_config, tool_context
+      self, credential_service, auth_config, callback_context
   ):
-    """Test loading credential with empty credential key."""
-    # Set credential key to empty string
+    """Test that loading with an empty credential key returns None."""
+    # Set credential_key to empty string
     auth_config.credential_key = ""
 
-    # Save first to have something to load
-    await credential_service.save_credential(auth_config, tool_context)
-
-    # Load should work with empty key
-    result = await credential_service.load_credential(auth_config, tool_context)
-    assert result == auth_config.exchanged_auth_credential
+    # Try to load credential
+    result = await credential_service.load_credential(
+        auth_config, callback_context
+    )
+    assert result is None
 
   @pytest.mark.asyncio
   async def test_state_persistence_across_operations(
-      self, credential_service, auth_config, tool_context
+      self, credential_service, auth_config, callback_context
   ):
-    """Test that state persists correctly across multiple operations."""
-    # Initially, no credential should exist
-    result = await credential_service.load_credential(auth_config, tool_context)
-    assert result is None
+    """Test that state persists across multiple operations."""
+    # Save credential
+    await credential_service.save_credential(auth_config, callback_context)
 
-    # Save a credential
-    await credential_service.save_credential(auth_config, tool_context)
+    # Verify state contains the credential
+    assert auth_config.credential_key in callback_context.state
+    assert (
+        callback_context.state[auth_config.credential_key]
+        == auth_config.exchanged_auth_credential
+    )
 
-    # Verify it was saved
-    result = await credential_service.load_credential(auth_config, tool_context)
+    # Load credential
+    result = await credential_service.load_credential(
+        auth_config, callback_context
+    )
     assert result is not None
-    assert result == auth_config.exchanged_auth_credential
 
-    # Update and save again
+    # Verify state still contains the credential
+    assert auth_config.credential_key in callback_context.state
+    assert (
+        callback_context.state[auth_config.credential_key]
+        == auth_config.exchanged_auth_credential
+    )
+
+    # Update credential
     new_credential = AuthCredential(
         auth_type=AuthCredentialTypes.OAUTH2,
         oauth2=OAuth2Auth(
-            client_id="new_client_id",
-            client_secret="new_client_secret",
-            redirect_uri="https://new.com/callback",
+            client_id="updated_client_id",
+            client_secret="updated_client_secret",
+            redirect_uri="https://updated.com/callback",
         ),
     )
     auth_config.exchanged_auth_credential = new_credential
-    await credential_service.save_credential(auth_config, tool_context)
 
-    # Verify the update persisted
-    result = await credential_service.load_credential(auth_config, tool_context)
-    assert result is not None
-    assert result.oauth2.client_id == "new_client_id"
+    # Save updated credential
+    await credential_service.save_credential(auth_config, callback_context)
+
+    # Verify state was updated
+    assert callback_context.state[auth_config.credential_key] == new_credential
 
   @pytest.mark.asyncio
   async def test_credential_key_uniqueness(
-      self, credential_service, oauth2_auth_scheme, tool_context
+      self, credential_service, oauth2_auth_scheme, callback_context
   ):
-    """Test that different credential keys create separate storage slots."""
-    # Create credentials with same content but different keys
-    credential = AuthCredential(
+    """Test that different credential keys store different credentials."""
+    # Create credentials with different keys
+    cred1 = AuthCredential(
         auth_type=AuthCredentialTypes.OAUTH2,
         oauth2=OAuth2Auth(
-            client_id="same_client",
-            client_secret="same_secret",
-            redirect_uri="https://same.com/callback",
+            client_id="client1",
+            client_secret="secret1",
+            redirect_uri="https://example1.com/callback",
         ),
     )
 
-    config_key1 = AuthConfig(
+    cred2 = AuthCredential(
+        auth_type=AuthCredentialTypes.OAUTH2,
+        oauth2=OAuth2Auth(
+            client_id="client2",
+            client_secret="secret2",
+            redirect_uri="https://example2.com/callback",
+        ),
+    )
+
+    auth_config1 = AuthConfig(
         auth_scheme=oauth2_auth_scheme,
-        raw_auth_credential=credential,
-        exchanged_auth_credential=credential,
+        raw_auth_credential=cred1,
+        exchanged_auth_credential=cred1,
         credential_key="unique_key_1",
     )
 
-    config_key2 = AuthConfig(
+    auth_config2 = AuthConfig(
         auth_scheme=oauth2_auth_scheme,
-        raw_auth_credential=credential,
-        exchanged_auth_credential=credential,
+        raw_auth_credential=cred2,
+        exchanged_auth_credential=cred2,
         credential_key="unique_key_2",
     )
 
-    # Save credential with first key
-    await credential_service.save_credential(config_key1, tool_context)
+    # Save both credentials
+    await credential_service.save_credential(auth_config1, callback_context)
+    await credential_service.save_credential(auth_config2, callback_context)
 
-    # Verify it's stored under first key
-    result1 = await credential_service.load_credential(
-        config_key1, tool_context
+    # Verify both exist in state with different keys
+    assert "unique_key_1" in callback_context.state
+    assert "unique_key_2" in callback_context.state
+    assert (
+        callback_context.state["unique_key_1"]
+        != callback_context.state["unique_key_2"]
     )
-    assert result1 is not None
 
-    # Verify it's not accessible under second key
-    result2 = await credential_service.load_credential(
-        config_key2, tool_context
-    )
-    assert result2 is None
-
-    # Save under second key
-    await credential_service.save_credential(config_key2, tool_context)
-
-    # Now both should be accessible
+    # Load and verify both credentials
     result1 = await credential_service.load_credential(
-        config_key1, tool_context
+        auth_config1, callback_context
     )
     result2 = await credential_service.load_credential(
-        config_key2, tool_context
+        auth_config2, callback_context
     )
+
     assert result1 is not None
     assert result2 is not None
-    assert result1 == result2  # Same credential content
+    assert result1.oauth2.client_id == "client1"
+    assert result2.oauth2.client_id == "client2"
 
-  def test_direct_state_access(
-      self, credential_service, auth_config, tool_context
+  @pytest.mark.asyncio
+  async def test_direct_state_access(
+      self, credential_service, auth_config, callback_context
   ):
-    """Test that the service correctly uses tool_context.state for storage."""
-    # Verify that the state starts empty
-    assert len(tool_context.state) == 0
+    """Test that the service properly accesses callback context state."""
+    # Directly set a value in state
+    test_credential = AuthCredential(
+        auth_type=AuthCredentialTypes.OAUTH2,
+        oauth2=OAuth2Auth(
+            client_id="direct_client_id",
+            client_secret="direct_client_secret",
+            redirect_uri="https://direct.com/callback",
+        ),
+    )
+    callback_context.state[auth_config.credential_key] = test_credential
 
-    # Save a credential (this is async but we're testing the state directly)
-    credential_key = auth_config.credential_key
-    test_credential = auth_config.exchanged_auth_credential
-
-    # Directly set the state to simulate save_credential behavior
-    tool_context.state[credential_key] = test_credential
-
-    # Verify the credential is in the state
-    assert credential_key in tool_context.state
-    assert tool_context.state[credential_key] == test_credential
-
-    # Verify we can retrieve it using the get method (simulating load_credential)
-    retrieved = tool_context.state.get(credential_key)
-    assert retrieved == test_credential
+    # Load using the service
+    result = await credential_service.load_credential(
+        auth_config, callback_context
+    )
+    assert result == test_credential

tests/unittests/auth/test_credential_manager.py

@@ -16,19 +16,13 @@ from unittest.mock import AsyncMock
 from unittest.mock import Mock
 from unittest.mock import patch
 
-from fastapi.openapi.models import HTTPBearer
-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 HttpAuth
-from google.adk.auth.auth_credential import HttpCredentials
 from google.adk.auth.auth_credential import OAuth2Auth
 from google.adk.auth.auth_credential import ServiceAccount
 from google.adk.auth.auth_credential import ServiceAccountCredential
+from google.adk.auth.auth_schemes import AuthScheme
 from google.adk.auth.auth_schemes import AuthSchemeType
-from google.adk.auth.auth_schemes import OpenIdConnectWithConfig
 from google.adk.auth.auth_tool import AuthConfig
 from google.adk.auth.credential_manager import CredentialManager
 import pytest
@@ -47,13 +41,13 @@ class TestCredentialManager:
   async def test_request_credential(self):
     """Test request_credential method."""
     auth_config = Mock(spec=AuthConfig)
-    tool_context = Mock()
-    tool_context.request_credential = Mock()
+    callback_context = Mock()
+    callback_context.request_credential = Mock()
 
     manager = CredentialManager(auth_config)
-    await manager.request_credential(tool_context)
+    await manager.request_credential(callback_context)
 
-    tool_context.request_credential.assert_called_once_with(auth_config)
+    callback_context.request_credential.assert_called_once_with(auth_config)
 
   @pytest.mark.asyncio
   async def test_load_auth_credentials_success(self):
@@ -67,7 +61,7 @@ class TestCredentialManager:
     mock_credential = Mock(spec=AuthCredential)
     mock_credential.auth_type = AuthCredentialTypes.API_KEY
 
-    tool_context = Mock()
+    callback_context = Mock()
 
     manager = CredentialManager(auth_config)
 
@@ -84,17 +78,17 @@ class TestCredentialManager:
     )
     manager._save_credential = AsyncMock()
 
-    result = await manager.get_auth_credential(tool_context)
+    result = await manager.get_auth_credential(callback_context)
 
     # Verify all methods were called
     manager._validate_credential.assert_called_once()
     manager._is_credential_ready.assert_called_once()
-    manager._load_existing_credential.assert_called_once_with(tool_context)
-    manager._load_from_auth_response.assert_called_once_with(tool_context)
+    manager._load_existing_credential.assert_called_once_with(callback_context)
+    manager._load_from_auth_response.assert_called_once_with(callback_context)
     manager._exchange_credential.assert_called_once_with(mock_credential)
     manager._refresh_credential.assert_called_once_with(mock_credential)
     manager._save_credential.assert_called_once_with(
-        tool_context, mock_credential
+        callback_context, mock_credential
     )
 
     assert result == mock_credential
@@ -106,7 +100,7 @@ class TestCredentialManager:
     auth_config.raw_auth_credential = None
     auth_config.exchanged_auth_credential = None
 
-    tool_context = Mock()
+    callback_context = Mock()
 
     manager = CredentialManager(auth_config)
 
@@ -115,20 +109,14 @@ class TestCredentialManager:
     manager._is_credential_ready = Mock(return_value=False)
     manager._load_existing_credential = AsyncMock(return_value=None)
     manager._load_from_auth_response = AsyncMock(return_value=None)
-    manager._exchange_credential = AsyncMock()
-    manager._refresh_credential = AsyncMock()
-    manager._save_credential = AsyncMock()
 
-    result = await manager.get_auth_credential(tool_context)
+    result = await manager.get_auth_credential(callback_context)
 
     # Verify methods were called but no credential returned
     manager._validate_credential.assert_called_once()
     manager._is_credential_ready.assert_called_once()
-    manager._load_existing_credential.assert_called_once_with(tool_context)
-    manager._load_from_auth_response.assert_called_once_with(tool_context)
-    manager._exchange_credential.assert_not_called()
-    manager._refresh_credential.assert_not_called()
-    manager._save_credential.assert_not_called()
+    manager._load_existing_credential.assert_called_once_with(callback_context)
+    manager._load_from_auth_response.assert_called_once_with(callback_context)
 
     assert result is None
 
@@ -139,12 +127,12 @@ class TestCredentialManager:
     mock_credential = Mock(spec=AuthCredential)
     auth_config.exchanged_auth_credential = mock_credential
 
-    tool_context = Mock()
+    callback_context = Mock()
 
     manager = CredentialManager(auth_config)
     manager._load_from_credential_service = AsyncMock(return_value=None)
 
-    result = await manager._load_existing_credential(tool_context)
+    result = await manager._load_existing_credential(callback_context)
 
     assert result == mock_credential
 
@@ -156,21 +144,23 @@ class TestCredentialManager:
 
     mock_credential = Mock(spec=AuthCredential)
 
-    tool_context = Mock()
+    callback_context = Mock()
 
     manager = CredentialManager(auth_config)
     manager._load_from_credential_service = AsyncMock(
         return_value=mock_credential
     )
 
-    result = await manager._load_existing_credential(tool_context)
+    result = await manager._load_existing_credential(callback_context)
 
-    manager._load_from_credential_service.assert_called_once_with(tool_context)
+    manager._load_from_credential_service.assert_called_once_with(
+        callback_context
+    )
     assert result == mock_credential
 
   @pytest.mark.asyncio
   async def test_load_from_credential_service_with_service(self):
-    """Test _load_from_credential_service from tool context when credential service is available."""
+    """Test _load_from_credential_service from callback context when credential service is available."""
     auth_config = Mock(spec=AuthConfig)
 
     mock_credential = Mock(spec=AuthCredential)
@@ -183,14 +173,14 @@ class TestCredentialManager:
     invocation_context = Mock()
     invocation_context.credential_service = credential_service
 
-    tool_context = Mock()
-    tool_context._invocation_context = invocation_context
+    callback_context = Mock()
+    callback_context._invocation_context = invocation_context
 
     manager = CredentialManager(auth_config)
-    result = await manager._load_from_credential_service(tool_context)
+    result = await manager._load_from_credential_service(callback_context)
 
     credential_service.load_credential.assert_called_once_with(
-        auth_config, tool_context
+        auth_config, callback_context
     )
     assert result == mock_credential
 
@@ -203,11 +193,11 @@ class TestCredentialManager:
     invocation_context = Mock()
     invocation_context.credential_service = None
 
-    tool_context = Mock()
-    tool_context._invocation_context = invocation_context
+    callback_context = Mock()
+    callback_context._invocation_context = invocation_context
 
     manager = CredentialManager(auth_config)
-    result = await manager._load_from_credential_service(tool_context)
+    result = await manager._load_from_credential_service(callback_context)
 
     assert result is None
 
@@ -224,14 +214,14 @@ class TestCredentialManager:
     invocation_context = Mock()
     invocation_context.credential_service = credential_service
 
-    tool_context = Mock()
-    tool_context._invocation_context = invocation_context
+    callback_context = Mock()
+    callback_context._invocation_context = invocation_context
 
     manager = CredentialManager(auth_config)
-    await manager._save_credential(tool_context, mock_credential)
+    await manager._save_credential(callback_context, mock_credential)
 
     credential_service.save_credential.assert_called_once_with(
-        auth_config, tool_context
+        auth_config, callback_context
     )
     assert auth_config.exchanged_auth_credential == mock_credential
 
@@ -246,11 +236,11 @@ class TestCredentialManager:
     invocation_context = Mock()
     invocation_context.credential_service = None
 
-    tool_context = Mock()
-    tool_context._invocation_context = invocation_context
+    callback_context = Mock()
+    callback_context._invocation_context = invocation_context
 
     manager = CredentialManager(auth_config)
-    await manager._save_credential(tool_context, mock_credential)
+    await manager._save_credential(callback_context, mock_credential)
 
     # Should not raise an error, and credential should not be set in auth_config
     # when there's no credential service (according to implementation)
@@ -383,58 +373,63 @@ class TestCredentialManager:
     auth_config.auth_scheme = auth_scheme
 
     manager = CredentialManager(auth_config)
-    await manager._validate_credential()
 
-    # Should return without error for non-OAuth2/OpenID schemes
+    # Should not raise an error for non-OAuth schemes
+    await manager._validate_credential()
 
   @pytest.mark.asyncio
   async def test_validate_credential_oauth2_missing_oauth2_field(self):
     """Test _validate_credential with OAuth2 credential missing oauth2 field."""
-    auth_scheme = Mock()
-    auth_scheme.type_ = AuthSchemeType.oauth2
-
     mock_raw_credential = Mock(spec=AuthCredential)
     mock_raw_credential.auth_type = AuthCredentialTypes.OAUTH2
     mock_raw_credential.oauth2 = None
 
     auth_config = Mock(spec=AuthConfig)
     auth_config.raw_auth_credential = mock_raw_credential
-    auth_config.auth_scheme = auth_scheme
-
-    manager = CredentialManager(auth_config)
-
-    with pytest.raises(
-        ValueError, match="auth_config.raw_credential.oauth2 required"
-    ):
-      await manager._validate_credential()
-
-  @pytest.mark.asyncio
-  async def test_exchange_credentials_service_account(self):
-    """Test _exchange_credential with service account credential (no exchanger available)."""
-    mock_raw_credential = Mock(spec=AuthCredential)
-    mock_raw_credential.auth_type = AuthCredentialTypes.SERVICE_ACCOUNT
-
-    auth_config = Mock(spec=AuthConfig)
     auth_config.auth_scheme = Mock()
 
     manager = CredentialManager(auth_config)
 
-    # Mock the exchanger registry to return None (no exchanger available)
+    with pytest.raises(ValueError, match="oauth2 required for credential type"):
+      await manager._validate_credential()
+
+  @pytest.mark.asyncio
+  async def test_exchange_credentials_service_account(self):
+    """Test _exchange_credential with service account credential."""
+    mock_service_account = Mock(spec=ServiceAccount)
+    mock_credential = Mock(spec=AuthCredential)
+    mock_credential.auth_type = AuthCredentialTypes.SERVICE_ACCOUNT
+
+    auth_config = Mock(spec=AuthConfig)
+    auth_config.auth_scheme = Mock()
+
+    # Mock exchanger
+    mock_exchanger = Mock()
+    mock_exchanger.exchange = AsyncMock(return_value=mock_credential)
+
+    manager = CredentialManager(auth_config)
+
+    # Mock the exchanger registry to return our mock exchanger
     with patch.object(
-        manager._exchanger_registry, "get_exchanger", return_value=None
+        manager._exchanger_registry,
+        "get_exchanger",
+        return_value=mock_exchanger,
     ):
       result, was_exchanged = await manager._exchange_credential(
-          mock_raw_credential
+          mock_credential
       )
 
-      assert result == mock_raw_credential
-      assert was_exchanged is False
+    mock_exchanger.exchange.assert_called_once_with(
+        mock_credential, auth_config.auth_scheme
+    )
+    assert result == mock_credential
+    assert was_exchanged is True
 
   @pytest.mark.asyncio
   async def test_exchange_credential_no_exchanger(self):
     """Test _exchange_credential with credential that has no exchanger."""
-    mock_raw_credential = Mock(spec=AuthCredential)
-    mock_raw_credential.auth_type = AuthCredentialTypes.API_KEY
+    mock_credential = Mock(spec=AuthCredential)
+    mock_credential.auth_type = AuthCredentialTypes.API_KEY
 
     auth_config = Mock(spec=AuthConfig)
 
@@ -442,55 +437,50 @@ class TestCredentialManager:
 
     # Mock the exchanger registry to return None (no exchanger available)
     with patch.object(
-        manager._exchanger_registry, "get_exchanger", return_value=None
+        manager._exchanger_registry,
+        "get_exchanger",
+        return_value=None,
     ):
       result, was_exchanged = await manager._exchange_credential(
-          mock_raw_credential
+          mock_credential
       )
 
-      assert result == mock_raw_credential
-      assert was_exchanged is False
+    assert result == mock_credential
+    assert was_exchanged is False
 
 
-# Test fixtures
 @pytest.fixture
 def oauth2_auth_scheme():
-  """Create an OAuth2 auth scheme for testing."""
-  flows = OAuthFlows(
-      authorizationCode=OAuthFlowAuthorizationCode(
-          authorizationUrl="https://example.com/oauth2/authorize",
-          tokenUrl="https://example.com/oauth2/token",
-          scopes={"read": "Read access", "write": "Write access"},
-      )
-  )
-  return OAuth2(flows=flows)
+  """OAuth2 auth scheme for testing."""
+  auth_scheme = Mock(spec=AuthScheme)
+  auth_scheme.type_ = AuthSchemeType.oauth2
+  return auth_scheme
 
 
 @pytest.fixture
 def openid_auth_scheme():
-  """Create an OpenID Connect auth scheme for testing."""
-  return OpenIdConnectWithConfig(
-      type_="openIdConnect",
-      authorization_endpoint="https://example.com/auth",
-      token_endpoint="https://example.com/token",
-      scopes=["openid", "profile"],
-  )
+  """OpenID Connect auth scheme for testing."""
+  auth_scheme = Mock(spec=AuthScheme)
+  auth_scheme.type_ = AuthSchemeType.openIdConnect
+  return auth_scheme
 
 
 @pytest.fixture
 def bearer_auth_scheme():
-  """Create a Bearer auth scheme for testing."""
-  return HTTPBearer(bearerFormat="JWT")
+  """Bearer auth scheme for testing."""
+  auth_scheme = Mock(spec=AuthScheme)
+  auth_scheme.type_ = AuthSchemeType.http
+  return auth_scheme
 
 
 @pytest.fixture
 def oauth2_credential():
-  """Create OAuth2 credentials for testing."""
+  """OAuth2 credential for testing."""
   return AuthCredential(
       auth_type=AuthCredentialTypes.OAUTH2,
       oauth2=OAuth2Auth(
-          client_id="mock_client_id",
-          client_secret="mock_client_secret",
+          client_id="test_client_id",
+          client_secret="test_client_secret",
           redirect_uri="https://example.com/callback",
       ),
   )
@@ -498,26 +488,27 @@ def oauth2_credential():
 
 @pytest.fixture
 def service_account_credential():
-  """Create service account credentials for testing."""
+  """Service account credential for testing."""
   return AuthCredential(
       auth_type=AuthCredentialTypes.SERVICE_ACCOUNT,
       service_account=ServiceAccount(
           service_account_credential=ServiceAccountCredential(
-              type="service_account",
-              project_id="test-project",
-              private_key_id="key-id",
+              type_="service_account",
+              project_id="test_project",
+              private_key_id="test_key_id",
               private_key=(
-                  "-----BEGIN PRIVATE KEY-----\ntest\n-----END PRIVATE"
+                  "-----BEGIN PRIVATE KEY-----\ntest_key\n-----END PRIVATE"
                   " KEY-----\n"
               ),
-              client_email="test@test-project.iam.gserviceaccount.com",
-              client_id="123456789",
+              client_email="test@test.iam.gserviceaccount.com",
+              client_id="test_client_id",
               auth_uri="https://accounts.google.com/o/oauth2/auth",
               token_uri="https://oauth2.googleapis.com/token",
               auth_provider_x509_cert_url=(
                   "https://www.googleapis.com/oauth2/v1/certs"
               ),
-              client_x509_cert_url="https://www.googleapis.com/robot/v1/metadata/x509/test%40test-project.iam.gserviceaccount.com",
+              client_x509_cert_url="https://www.googleapis.com/robot/v1/metadata/x509/test%40test.iam.gserviceaccount.com",
+              universe_domain="googleapis.com",
           ),
           scopes=["https://www.googleapis.com/auth/cloud-platform"],
       ),
@@ -526,20 +517,17 @@ def service_account_credential():
 
 @pytest.fixture
 def api_key_credential():
-  """Create API key credentials for testing."""
+  """API key credential for testing."""
   return AuthCredential(
       auth_type=AuthCredentialTypes.API_KEY,
-      api_key="test-api-key",
+      api_key="test_api_key",
   )
 
 
 @pytest.fixture
 def http_bearer_credential():
-  """Create HTTP Bearer credentials for testing."""
+  """HTTP bearer credential for testing."""
   return AuthCredential(
       auth_type=AuthCredentialTypes.HTTP,
-      http=HttpAuth(
-          scheme="bearer",
-          credentials=HttpCredentials(token="bearer-token"),
-      ),
+      http=Mock(),
   )