encounter/adk-python
0
0
Fork 0
mirror of https://github.com/encounter/adk-python.git synced 2026-07-09 18:19:28 -07:00
Code Issues Packages Projects Releases Wiki Activity

chore: Remove service account support

Browse Source 
given it was not correctly supported.

PiperOrigin-RevId: 773137317
This commit is contained in:
Xiang (Sean) Zhou
2025-06-18 18:19:05 -07:00
committed by Copybara-Service
parent 913d771d6d
commit 9a1115c504
11 changed files with 15 additions and 760 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/google/adk/auth/auth_credential.py
-1
View File
@@ -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
src/google/adk/auth/credential_manager.py
+1 -5
View File
@@ -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()
src/google/adk/auth/exchanger/__init__.py
-2
View File
@@ -15,9 +15,7 @@
"""Credential exchanger module."""
from .base_credential_exchanger import BaseCredentialExchanger
from .service_account_credential_exchanger import ServiceAccountCredentialExchanger
__all__ = [
"BaseCredentialExchanger",
"ServiceAccountCredentialExchanger",
]
src/google/adk/auth/exchanger/service_account_credential_exchanger.py
-104
View File
@@ -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
src/google/adk/auth/refresher/oauth2_credential_refresher.py
+2 -30
View File
@@ -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
src/google/adk/tools/mcp_tool/mcp_tool.py
+2 -8
View File
@@ -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
src/google/adk/tools/openapi_tool/openapi_spec_parser/tool_auth_handler.py
-1
View File
@@ -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(
tests/unittests/auth/exchanger/test_service_account_credential_exchanger.py
-433
View File
@@ -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)
tests/unittests/auth/refresher/test_oauth2_credential_refresher.py
-118
View File
@@ -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."""
tests/unittests/auth/test_credential_manager.py
+6 -20
View File
@@ -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):
tests/unittests/tools/mcp_tool/test_mcp_tool.py
+4 -38
View File
@@ -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