diff --git a/src/google/adk/auth/auth_credential.py b/src/google/adk/auth/auth_credential.py index 1009a50d..34d04dde 100644 --- a/src/google/adk/auth/auth_credential.py +++ b/src/google/adk/auth/auth_credential.py @@ -230,4 +230,3 @@ class AuthCredential(BaseModelWithConfig): http: Optional[HttpAuth] = None service_account: Optional[ServiceAccount] = None oauth2: Optional[OAuth2Auth] = None - google_oauth2_json: Optional[str] = None diff --git a/src/google/adk/auth/credential_manager.py b/src/google/adk/auth/credential_manager.py index 7471bdff..0dbf006a 100644 --- a/src/google/adk/auth/credential_manager.py +++ b/src/google/adk/auth/credential_manager.py @@ -76,11 +76,7 @@ class CredentialManager: self._refresher_registry = CredentialRefresherRegistry() # Register default exchangers and refreshers - from .exchanger.service_account_credential_exchanger import ServiceAccountCredentialExchanger - - self._exchanger_registry.register( - AuthCredentialTypes.SERVICE_ACCOUNT, ServiceAccountCredentialExchanger() - ) + # TODO: support service account credential exchanger from .refresher.oauth2_credential_refresher import OAuth2CredentialRefresher oauth2_refresher = OAuth2CredentialRefresher() diff --git a/src/google/adk/auth/exchanger/__init__.py b/src/google/adk/auth/exchanger/__init__.py index 4226ae71..3b0fbb24 100644 --- a/src/google/adk/auth/exchanger/__init__.py +++ b/src/google/adk/auth/exchanger/__init__.py @@ -15,9 +15,7 @@ """Credential exchanger module.""" from .base_credential_exchanger import BaseCredentialExchanger -from .service_account_credential_exchanger import ServiceAccountCredentialExchanger __all__ = [ "BaseCredentialExchanger", - "ServiceAccountCredentialExchanger", ] diff --git a/src/google/adk/auth/exchanger/service_account_credential_exchanger.py b/src/google/adk/auth/exchanger/service_account_credential_exchanger.py deleted file mode 100644 index 415081ca..00000000 --- a/src/google/adk/auth/exchanger/service_account_credential_exchanger.py +++ /dev/null @@ -1,104 +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. - -"""Credential fetcher for Google Service Account.""" - -from __future__ import annotations - -from typing import Optional - -import google.auth -from google.auth.transport.requests import Request -from google.oauth2 import service_account -from typing_extensions import override - -from ...utils.feature_decorator import experimental -from ..auth_credential import AuthCredential -from ..auth_credential import AuthCredentialTypes -from ..auth_schemes import AuthScheme -from .base_credential_exchanger import BaseCredentialExchanger - - -@experimental -class ServiceAccountCredentialExchanger(BaseCredentialExchanger): - """Exchanges Google Service Account credentials for an access token. - - Uses the default service credential if `use_default_credential = True`. - Otherwise, uses the service account credential provided in the auth - credential. - """ - - @override - async def exchange( - self, - auth_credential: AuthCredential, - auth_scheme: Optional[AuthScheme] = None, - ) -> AuthCredential: - """Exchanges the service account auth credential for an access token. - - If the AuthCredential contains a service account credential, it will be used - to exchange for an access token. Otherwise, if use_default_credential is True, - the default application credential will be used for exchanging an access token. - - Args: - auth_scheme: The authentication scheme. - auth_credential: The credential to exchange. - - Returns: - An AuthCredential in OAUTH2 format, containing the exchanged credential JSON. - - Raises: - ValueError: If service account credentials are missing or invalid. - Exception: If credential exchange or refresh fails. - """ - if auth_credential is None: - raise ValueError("Credential cannot be None.") - - if auth_credential.auth_type != AuthCredentialTypes.SERVICE_ACCOUNT: - raise ValueError("Credential is not a service account credential.") - - if auth_credential.service_account is None: - raise ValueError( - "Service account credentials are missing. Please provide them." - ) - - if ( - auth_credential.service_account.service_account_credential is None - and not auth_credential.service_account.use_default_credential - ): - raise ValueError( - "Service account credentials are invalid. Please set the" - " service_account_credential field or set `use_default_credential =" - " True` to use application default credential in a hosted service" - " like Google Cloud Run." - ) - - try: - if auth_credential.service_account.use_default_credential: - credentials, _ = google.auth.default() - else: - config = auth_credential.service_account - credentials = service_account.Credentials.from_service_account_info( - config.service_account_credential.model_dump(), scopes=config.scopes - ) - - # Refresh credentials to ensure we have a valid access token - credentials.refresh(Request()) - - return AuthCredential( - auth_type=AuthCredentialTypes.OAUTH2, - google_oauth2_json=credentials.to_json(), - ) - except Exception as e: - raise ValueError(f"Failed to exchange service account token: {e}") from e diff --git a/src/google/adk/auth/refresher/oauth2_credential_refresher.py b/src/google/adk/auth/refresher/oauth2_credential_refresher.py index 2d0a8b67..4c19520c 100644 --- a/src/google/adk/auth/refresher/oauth2_credential_refresher.py +++ b/src/google/adk/auth/refresher/oauth2_credential_refresher.py @@ -60,27 +60,12 @@ class OAuth2CredentialRefresher(BaseCredentialRefresher): Returns: True if the credential needs to be refreshed, False otherwise. """ - # Handle Google OAuth2 credentials (from service account exchange) - if auth_credential.google_oauth2_json: - try: - google_credential = Credentials.from_authorized_user_info( - json.loads(auth_credential.google_oauth2_json) - ) - return google_credential.expired and bool( - google_credential.refresh_token - ) - except Exception as e: - logger.warning("Failed to parse Google OAuth2 JSON credential: %s", e) - return False # Handle regular OAuth2 credentials - elif auth_credential.oauth2 and auth_scheme: + if auth_credential.oauth2: if not AUTHLIB_AVIALABLE: return False - if not auth_credential.oauth2: - return False - return OAuth2Token({ "expires_at": auth_credential.oauth2.expires_at, "expires_in": auth_credential.oauth2.expires_in, @@ -105,22 +90,9 @@ class OAuth2CredentialRefresher(BaseCredentialRefresher): The refreshed credential. """ - # Handle Google OAuth2 credentials (from service account exchange) - if auth_credential.google_oauth2_json: - try: - google_credential = Credentials.from_authorized_user_info( - json.loads(auth_credential.google_oauth2_json) - ) - if google_credential.expired and google_credential.refresh_token: - google_credential.refresh(Request()) - auth_credential.google_oauth2_json = google_credential.to_json() - logger.info("Successfully refreshed Google OAuth2 JSON credential") - except Exception as e: - # TODO reconsider whether we should raise error when refresh failed. - logger.error("Failed to refresh Google OAuth2 JSON credential: %s", e) # Handle regular OAuth2 credentials - elif auth_credential.oauth2 and auth_scheme: + if auth_credential.oauth2 and auth_scheme: if not AUTHLIB_AVIALABLE: return auth_credential diff --git a/src/google/adk/tools/mcp_tool/mcp_tool.py b/src/google/adk/tools/mcp_tool/mcp_tool.py index 24998c92..310fc48f 100644 --- a/src/google/adk/tools/mcp_tool/mcp_tool.py +++ b/src/google/adk/tools/mcp_tool/mcp_tool.py @@ -138,11 +138,6 @@ class MCPTool(BaseAuthenticatedTool): if credential: if credential.oauth2: headers = {"Authorization": f"Bearer {credential.oauth2.access_token}"} - elif credential.google_oauth2_json: - google_credential = Credentials.from_authorized_user_info( - json.loads(credential.google_oauth2_json) - ) - headers = {"Authorization": f"Bearer {google_credential.token}"} elif credential.http: # Handle HTTP authentication schemes if ( @@ -178,10 +173,9 @@ class MCPTool(BaseAuthenticatedTool): headers = {"X-API-Key": credential.api_key} elif credential.service_account: # Service accounts should be exchanged for access tokens before reaching this point - # If we reach here, we can try to use google_oauth2_json or log a warning logger.warning( - "Service account credentials should be exchanged for access" - " tokens before MCP session creation" + "Service account credentials should be exchanged before MCP" + " session creation" ) return headers diff --git a/src/google/adk/tools/openapi_tool/openapi_spec_parser/tool_auth_handler.py b/src/google/adk/tools/openapi_tool/openapi_spec_parser/tool_auth_handler.py index 08e535d2..74166b00 100644 --- a/src/google/adk/tools/openapi_tool/openapi_spec_parser/tool_auth_handler.py +++ b/src/google/adk/tools/openapi_tool/openapi_spec_parser/tool_auth_handler.py @@ -233,7 +233,6 @@ class ToolAuthHandler: AuthCredentialTypes.OPEN_ID_CONNECT, ) and not credential.oauth2.access_token - and not credential.google_oauth2_json ) async def prepare_auth_credentials( diff --git a/tests/unittests/auth/exchanger/test_service_account_credential_exchanger.py b/tests/unittests/auth/exchanger/test_service_account_credential_exchanger.py deleted file mode 100644 index 195e143d..00000000 --- a/tests/unittests/auth/exchanger/test_service_account_credential_exchanger.py +++ /dev/null @@ -1,433 +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. - -"""Unit tests for the ServiceAccountCredentialExchanger.""" - -from unittest.mock import MagicMock -from unittest.mock import patch - -from fastapi.openapi.models import HTTPBearer -from google.adk.auth.auth_credential import AuthCredential -from google.adk.auth.auth_credential import AuthCredentialTypes -from google.adk.auth.auth_credential import ServiceAccount -from google.adk.auth.auth_credential import ServiceAccountCredential -from google.adk.auth.exchanger.service_account_credential_exchanger import ServiceAccountCredentialExchanger -import pytest - - -class TestServiceAccountCredentialExchanger: - """Test cases for ServiceAccountCredentialExchanger.""" - - def test_exchange_with_valid_credential(self): - """Test successful exchange with valid service account credential.""" - credential = 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", - private_key=( - "-----BEGIN PRIVATE KEY-----\nMOCK_KEY\n-----END PRIVATE" - " KEY-----" - ), - client_email="test@test-project.iam.gserviceaccount.com", - client_id="12345", - 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", - universe_domain="googleapis.com", - ), - scopes=["https://www.googleapis.com/auth/cloud-platform"], - ), - ) - - auth_scheme = HTTPBearer() - exchanger = ServiceAccountCredentialExchanger() - - # This should not raise an exception - assert exchanger is not None - - @pytest.mark.asyncio - async def test_exchange_invalid_credential_type(self): - """Test exchange with invalid credential type raises ValueError.""" - credential = AuthCredential( - auth_type=AuthCredentialTypes.API_KEY, - api_key="test-key", - ) - - auth_scheme = HTTPBearer() - exchanger = ServiceAccountCredentialExchanger() - - with pytest.raises( - ValueError, match="Credential is not a service account credential" - ): - await exchanger.exchange(credential, auth_scheme) - - @pytest.mark.asyncio - @patch( - "google.adk.auth.exchanger.service_account_credential_exchanger.service_account.Credentials.from_service_account_info" - ) - @patch( - "google.adk.auth.exchanger.service_account_credential_exchanger.Request" - ) - async def test_exchange_with_explicit_credentials_success( - self, mock_request_class, mock_from_service_account_info - ): - """Test successful exchange with explicit service account credentials.""" - # Setup mocks - mock_request = MagicMock() - mock_request_class.return_value = mock_request - - mock_credentials = MagicMock() - mock_credentials.token = "mock_access_token" - mock_credentials.to_json.return_value = ( - '{"token": "mock_access_token", "type": "authorized_user"}' - ) - mock_from_service_account_info.return_value = mock_credentials - - # Create test credential - service_account_cred = ServiceAccountCredential( - type_="service_account", - project_id="test-project", - private_key_id="key-id", - private_key=( - "-----BEGIN PRIVATE KEY-----\nMOCK_KEY\n-----END PRIVATE KEY-----" - ), - client_email="test@test-project.iam.gserviceaccount.com", - client_id="12345", - 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", - universe_domain="googleapis.com", - ) - - credential = AuthCredential( - auth_type=AuthCredentialTypes.SERVICE_ACCOUNT, - service_account=ServiceAccount( - service_account_credential=service_account_cred, - scopes=["https://www.googleapis.com/auth/cloud-platform"], - ), - ) - - auth_scheme = HTTPBearer() - exchanger = ServiceAccountCredentialExchanger() - result = await exchanger.exchange(credential, auth_scheme) - - # Verify the result - assert result.auth_type == AuthCredentialTypes.OAUTH2 - assert result.google_oauth2_json is not None - # Verify that google_oauth2_json contains the token - import json - - exchanged_creds = json.loads(result.google_oauth2_json) - assert exchanged_creds.get( - "token" - ) == "mock_access_token" or "mock_access_token" in str(exchanged_creds) - - # Verify mocks were called correctly - mock_from_service_account_info.assert_called_once_with( - service_account_cred.model_dump(), - scopes=["https://www.googleapis.com/auth/cloud-platform"], - ) - mock_credentials.refresh.assert_called_once_with(mock_request) - - @pytest.mark.asyncio - @patch( - "google.adk.auth.exchanger.service_account_credential_exchanger.google.auth.default" - ) - @patch( - "google.adk.auth.exchanger.service_account_credential_exchanger.Request" - ) - async def test_exchange_with_default_credentials_success( - self, mock_request_class, mock_google_auth_default - ): - """Test successful exchange with default application credentials.""" - # Setup mocks - mock_request = MagicMock() - mock_request_class.return_value = mock_request - - mock_credentials = MagicMock() - mock_credentials.token = "default_access_token" - mock_credentials.to_json.return_value = ( - '{"token": "default_access_token", "type": "authorized_user"}' - ) - mock_google_auth_default.return_value = (mock_credentials, "test-project") - - # Create test credential with use_default_credential=True - credential = AuthCredential( - auth_type=AuthCredentialTypes.SERVICE_ACCOUNT, - service_account=ServiceAccount( - use_default_credential=True, - scopes=["https://www.googleapis.com/auth/cloud-platform"], - ), - ) - - auth_scheme = HTTPBearer() - exchanger = ServiceAccountCredentialExchanger() - result = await exchanger.exchange(credential, auth_scheme) - - # Verify the result - assert result.auth_type == AuthCredentialTypes.OAUTH2 - assert result.google_oauth2_json is not None - # Verify that google_oauth2_json contains the token - import json - - exchanged_creds = json.loads(result.google_oauth2_json) - assert exchanged_creds.get( - "token" - ) == "default_access_token" or "default_access_token" in str( - exchanged_creds - ) - - # Verify mocks were called correctly - mock_google_auth_default.assert_called_once() - mock_credentials.refresh.assert_called_once_with(mock_request) - - @pytest.mark.asyncio - async def test_exchange_missing_service_account(self): - """Test exchange fails when service_account is None.""" - credential = AuthCredential( - auth_type=AuthCredentialTypes.SERVICE_ACCOUNT, - service_account=None, - ) - - auth_scheme = HTTPBearer() - exchanger = ServiceAccountCredentialExchanger() - - with pytest.raises( - ValueError, match="Service account credentials are missing" - ): - await exchanger.exchange(credential, auth_scheme) - - @pytest.mark.asyncio - async def test_exchange_missing_credentials_and_not_default(self): - """Test exchange fails when credentials are missing and use_default_credential is False.""" - credential = AuthCredential( - auth_type=AuthCredentialTypes.SERVICE_ACCOUNT, - service_account=ServiceAccount( - service_account_credential=None, - use_default_credential=False, - scopes=["https://www.googleapis.com/auth/cloud-platform"], - ), - ) - - auth_scheme = HTTPBearer() - exchanger = ServiceAccountCredentialExchanger() - - with pytest.raises( - ValueError, match="Service account credentials are invalid" - ): - await exchanger.exchange(credential, auth_scheme) - - @pytest.mark.asyncio - @patch( - "google.adk.auth.exchanger.service_account_credential_exchanger.service_account.Credentials.from_service_account_info" - ) - async def test_exchange_credential_creation_failure( - self, mock_from_service_account_info - ): - """Test exchange handles credential creation failure gracefully.""" - # Setup mock to raise exception - mock_from_service_account_info.side_effect = Exception( - "Invalid private key" - ) - - # Create test credential - service_account_cred = ServiceAccountCredential( - type_="service_account", - project_id="test-project", - private_key_id="key-id", - private_key="invalid-key", - client_email="test@test-project.iam.gserviceaccount.com", - client_id="12345", - 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", - universe_domain="googleapis.com", - ) - - credential = AuthCredential( - auth_type=AuthCredentialTypes.SERVICE_ACCOUNT, - service_account=ServiceAccount( - service_account_credential=service_account_cred, - scopes=["https://www.googleapis.com/auth/cloud-platform"], - ), - ) - - auth_scheme = HTTPBearer() - exchanger = ServiceAccountCredentialExchanger() - - with pytest.raises( - ValueError, match="Failed to exchange service account token" - ): - await exchanger.exchange(credential, auth_scheme) - - @pytest.mark.asyncio - @patch( - "google.adk.auth.exchanger.service_account_credential_exchanger.google.auth.default" - ) - async def test_exchange_default_credential_failure( - self, mock_google_auth_default - ): - """Test exchange handles default credential failure gracefully.""" - # Setup mock to raise exception - mock_google_auth_default.side_effect = Exception( - "No default credentials found" - ) - - # Create test credential with use_default_credential=True - credential = AuthCredential( - auth_type=AuthCredentialTypes.SERVICE_ACCOUNT, - service_account=ServiceAccount( - use_default_credential=True, - scopes=["https://www.googleapis.com/auth/cloud-platform"], - ), - ) - - auth_scheme = HTTPBearer() - exchanger = ServiceAccountCredentialExchanger() - - with pytest.raises( - ValueError, match="Failed to exchange service account token" - ): - await exchanger.exchange(credential, auth_scheme) - - @pytest.mark.asyncio - @patch( - "google.adk.auth.exchanger.service_account_credential_exchanger.service_account.Credentials.from_service_account_info" - ) - @patch( - "google.adk.auth.exchanger.service_account_credential_exchanger.Request" - ) - async def test_exchange_refresh_failure( - self, mock_request_class, mock_from_service_account_info - ): - """Test exchange handles credential refresh failure gracefully.""" - # Setup mocks - mock_request = MagicMock() - mock_request_class.return_value = mock_request - - mock_credentials = MagicMock() - mock_credentials.refresh.side_effect = Exception( - "Network error during refresh" - ) - mock_from_service_account_info.return_value = mock_credentials - - # Create test credential - service_account_cred = ServiceAccountCredential( - type_="service_account", - project_id="test-project", - private_key_id="key-id", - private_key=( - "-----BEGIN PRIVATE KEY-----\nMOCK_KEY\n-----END PRIVATE KEY-----" - ), - client_email="test@test-project.iam.gserviceaccount.com", - client_id="12345", - 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", - universe_domain="googleapis.com", - ) - - credential = AuthCredential( - auth_type=AuthCredentialTypes.SERVICE_ACCOUNT, - service_account=ServiceAccount( - service_account_credential=service_account_cred, - scopes=["https://www.googleapis.com/auth/cloud-platform"], - ), - ) - - auth_scheme = HTTPBearer() - exchanger = ServiceAccountCredentialExchanger() - - with pytest.raises( - ValueError, match="Failed to exchange service account token" - ): - await exchanger.exchange(credential, auth_scheme) - - @pytest.mark.asyncio - async def test_exchange_none_credential_in_constructor(self): - """Test that passing None credential raises appropriate error during exchange.""" - # This test verifies behavior when credential is None - auth_scheme = HTTPBearer() - exchanger = ServiceAccountCredentialExchanger() - - with pytest.raises(ValueError, match="Credential cannot be None"): - await exchanger.exchange(None, auth_scheme) - - @pytest.mark.asyncio - @patch( - "google.adk.auth.exchanger.service_account_credential_exchanger.google.auth.default" - ) - @patch( - "google.adk.auth.exchanger.service_account_credential_exchanger.Request" - ) - async def test_exchange_with_service_account_no_explicit_credentials( - self, mock_request_class, mock_google_auth_default - ): - """Test exchange with service account that has no explicit credentials uses default.""" - # Setup mocks - mock_request = MagicMock() - mock_request_class.return_value = mock_request - - mock_credentials = MagicMock() - mock_credentials.token = "default_access_token" - mock_credentials.to_json.return_value = ( - '{"token": "default_access_token", "type": "authorized_user"}' - ) - mock_google_auth_default.return_value = (mock_credentials, "test-project") - - # Create test credential with no explicit credentials but use_default_credential=True - credential = AuthCredential( - auth_type=AuthCredentialTypes.SERVICE_ACCOUNT, - service_account=ServiceAccount( - service_account_credential=None, - use_default_credential=True, - scopes=["https://www.googleapis.com/auth/cloud-platform"], - ), - ) - - auth_scheme = HTTPBearer() - exchanger = ServiceAccountCredentialExchanger() - result = await exchanger.exchange(credential, auth_scheme) - - # Verify the result - assert result.auth_type == AuthCredentialTypes.OAUTH2 - assert result.google_oauth2_json is not None - # Verify that google_oauth2_json contains the token - import json - - exchanged_creds = json.loads(result.google_oauth2_json) - assert exchanged_creds.get( - "token" - ) == "default_access_token" or "default_access_token" in str( - exchanged_creds - ) - - # Verify mocks were called correctly - mock_google_auth_default.assert_called_once() - mock_credentials.refresh.assert_called_once_with(mock_request) diff --git a/tests/unittests/auth/refresher/test_oauth2_credential_refresher.py b/tests/unittests/auth/refresher/test_oauth2_credential_refresher.py index b22bf2cc..3342fcb0 100644 --- a/tests/unittests/auth/refresher/test_oauth2_credential_refresher.py +++ b/tests/unittests/auth/refresher/test_oauth2_credential_refresher.py @@ -165,124 +165,6 @@ class TestOAuth2CredentialRefresher: 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.""" diff --git a/tests/unittests/auth/test_credential_manager.py b/tests/unittests/auth/test_credential_manager.py index 283e865a..8e3638dd 100644 --- a/tests/unittests/auth/test_credential_manager.py +++ b/tests/unittests/auth/test_credential_manager.py @@ -410,39 +410,25 @@ class TestCredentialManager: @pytest.mark.asyncio async def test_exchange_credentials_service_account(self): - """Test _exchange_credential with service account credential.""" + """Test _exchange_credential with service account credential (no exchanger available).""" mock_raw_credential = Mock(spec=AuthCredential) mock_raw_credential.auth_type = AuthCredentialTypes.SERVICE_ACCOUNT - mock_exchanged_credential = Mock(spec=AuthCredential) - auth_config = Mock(spec=AuthConfig) auth_config.auth_scheme = Mock() manager = CredentialManager(auth_config) - # Mock the exchanger that gets created during registration + # Mock the exchanger registry to return None (no exchanger available) with patch.object( - manager._exchanger_registry, "get_exchanger" - ) as mock_get_exchanger: - mock_exchanger = Mock() - mock_exchanger.exchange = AsyncMock( - return_value=mock_exchanged_credential - ) - mock_get_exchanger.return_value = mock_exchanger - + manager._exchanger_registry, "get_exchanger", return_value=None + ): result, was_exchanged = await manager._exchange_credential( mock_raw_credential ) - assert result == mock_exchanged_credential - assert was_exchanged is True - mock_get_exchanger.assert_called_once_with( - AuthCredentialTypes.SERVICE_ACCOUNT - ) - mock_exchanger.exchange.assert_called_once_with( - mock_raw_credential, auth_config.auth_scheme - ) + assert result == mock_raw_credential + assert was_exchanged is False @pytest.mark.asyncio async def test_exchange_credential_no_exchanger(self): diff --git a/tests/unittests/tools/mcp_tool/test_mcp_tool.py b/tests/unittests/tools/mcp_tool/test_mcp_tool.py index 4d9cffb4..d25a84ea 100644 --- a/tests/unittests/tools/mcp_tool/test_mcp_tool.py +++ b/tests/unittests/tools/mcp_tool/test_mcp_tool.py @@ -263,40 +263,6 @@ class TestMCPTool: assert headers == {"X-API-Key": "my_api_key"} - @pytest.mark.asyncio - @patch("google.adk.tools.mcp_tool.mcp_tool.json") - @patch("google.adk.tools.mcp_tool.mcp_tool.Credentials") - async def test_get_headers_google_oauth2_json( - self, mock_credentials, mock_json - ): - """Test header generation for Google OAuth2 JSON credentials.""" - tool = MCPTool( - mcp_tool=self.mock_mcp_tool, - mcp_session_manager=self.mock_session_manager, - ) - - # Mock the JSON parsing and Credentials creation - mock_json.loads.return_value = {"token": "google_token"} - mock_google_credential = Mock() - mock_google_credential.token = "google_access_token" - mock_credentials.from_authorized_user_info.return_value = ( - mock_google_credential - ) - - credential = AuthCredential( - auth_type=AuthCredentialTypes.OAUTH2, - google_oauth2_json='{"token": "google_token"}', - ) - - tool_context = Mock(spec=ToolContext) - headers = await tool._get_headers(tool_context, credential) - - assert headers == {"Authorization": "Bearer google_access_token"} - mock_json.loads.assert_called_once_with('{"token": "google_token"}') - mock_credentials.from_authorized_user_info.assert_called_once_with( - {"token": "google_token"} - ) - @pytest.mark.asyncio async def test_get_headers_no_credential(self): """Test header generation with no credentials.""" @@ -311,14 +277,14 @@ class TestMCPTool: assert headers is None @pytest.mark.asyncio - async def test_get_headers_service_account_no_json(self): - """Test header generation for service account credentials without google_oauth2_json.""" + async def test_get_headers_service_account(self): + """Test header generation for service account credentials.""" tool = MCPTool( mcp_tool=self.mock_mcp_tool, mcp_session_manager=self.mock_session_manager, ) - # Create service account credential without google_oauth2_json + # Create service account credential service_account = ServiceAccount(scopes=["test"]) credential = AuthCredential( auth_type=AuthCredentialTypes.SERVICE_ACCOUNT, @@ -328,7 +294,7 @@ class TestMCPTool: tool_context = Mock(spec=ToolContext) headers = await tool._get_headers(tool_context, credential) - # Should return None as no google_oauth2_json is provided + # Should return None as service account credentials are not supported for direct header generation assert headers is None @pytest.mark.asyncio