mirror of
https://github.com/encounter/adk-python.git
synced 2026-07-09 18:19:28 -07:00
feat: Support ID token exchange in ServiceAccountCredentialExchanger
Adds use_id_token and audience fields to ServiceAccount so that ServiceAccountCredentialExchanger can produce ID tokens instead of access tokens. This is required for authenticating to Cloud Run, Cloud Functions, and other Google Cloud services that verify caller identity. Close #4458 Co-authored-by: George Weale <gweale@google.com> PiperOrigin-RevId: 874630210
This commit is contained in:
George Weale
committed by
Copybara-Service
Copybara-Service
parent
c615757ba1
commit
7be90db24b
@@ -534,7 +534,9 @@ class TestMCPTool:
|
||||
)
|
||||
|
||||
# Create service account credential
|
||||
service_account = ServiceAccount(scopes=["test"])
|
||||
service_account = ServiceAccount(
|
||||
scopes=["test"], use_default_credential=True
|
||||
)
|
||||
credential = AuthCredential(
|
||||
auth_type=AuthCredentialTypes.SERVICE_ACCOUNT,
|
||||
service_account=service_account,
|
||||
|
||||
+253
-80
@@ -25,8 +25,23 @@ from google.adk.auth.auth_schemes import AuthSchemeType
|
||||
from google.adk.tools.openapi_tool.auth.credential_exchangers.base_credential_exchanger import AuthCredentialMissingError
|
||||
from google.adk.tools.openapi_tool.auth.credential_exchangers.service_account_exchanger import ServiceAccountCredentialExchanger
|
||||
import google.auth
|
||||
from google.auth import exceptions as google_auth_exceptions
|
||||
import pytest
|
||||
|
||||
_ACCESS_TOKEN_MONKEYPATCH_TARGET = (
|
||||
"google.adk.tools.openapi_tool.auth.credential_exchangers."
|
||||
"service_account_exchanger.service_account.Credentials."
|
||||
"from_service_account_info"
|
||||
)
|
||||
|
||||
_ID_TOKEN_MONKEYPATCH_TARGET = (
|
||||
"google.adk.tools.openapi_tool.auth.credential_exchangers."
|
||||
"service_account_exchanger.service_account.IDTokenCredentials."
|
||||
"from_service_account_info"
|
||||
)
|
||||
|
||||
_FETCH_ID_TOKEN_MONKEYPATCH_TARGET = "google.oauth2.id_token.fetch_id_token"
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def service_account_exchanger():
|
||||
@@ -41,50 +56,45 @@ def auth_scheme():
|
||||
return scheme
|
||||
|
||||
|
||||
def test_exchange_credential_success(
|
||||
service_account_exchanger, auth_scheme, monkeypatch
|
||||
@pytest.fixture
|
||||
def sa_credential():
|
||||
"""A minimal valid ServiceAccountCredential for testing."""
|
||||
return ServiceAccountCredential(
|
||||
type_="service_account",
|
||||
project_id="test_project_id",
|
||||
private_key_id="test_private_key_id",
|
||||
private_key="-----BEGIN PRIVATE KEY-----...",
|
||||
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"
|
||||
),
|
||||
universe_domain="googleapis.com",
|
||||
)
|
||||
|
||||
|
||||
_DEFAULT_SCOPES = ["https://www.googleapis.com/auth/cloud-platform"]
|
||||
|
||||
|
||||
# --- Access token exchange tests ---
|
||||
|
||||
|
||||
def test_exchange_access_token_with_explicit_credentials(
|
||||
service_account_exchanger, auth_scheme, sa_credential, monkeypatch
|
||||
):
|
||||
"""Test successful exchange of service account credentials."""
|
||||
mock_credentials = MagicMock()
|
||||
mock_credentials.token = "mock_access_token"
|
||||
mock_from_sa_info = MagicMock(return_value=mock_credentials)
|
||||
monkeypatch.setattr(_ACCESS_TOKEN_MONKEYPATCH_TARGET, mock_from_sa_info)
|
||||
|
||||
# Mock the from_service_account_info method
|
||||
mock_from_service_account_info = MagicMock(return_value=mock_credentials)
|
||||
target_path = (
|
||||
"google.adk.tools.openapi_tool.auth.credential_exchangers."
|
||||
"service_account_exchanger.service_account.Credentials."
|
||||
"from_service_account_info"
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
target_path,
|
||||
mock_from_service_account_info,
|
||||
)
|
||||
|
||||
# Mock the refresh method
|
||||
mock_credentials.refresh = MagicMock()
|
||||
|
||||
# Create a valid AuthCredential with service account info
|
||||
auth_credential = AuthCredential(
|
||||
auth_type=AuthCredentialTypes.SERVICE_ACCOUNT,
|
||||
service_account=ServiceAccount(
|
||||
service_account_credential=ServiceAccountCredential(
|
||||
type_="service_account",
|
||||
project_id="your_project_id",
|
||||
private_key_id="your_private_key_id",
|
||||
private_key="-----BEGIN PRIVATE KEY-----...",
|
||||
client_email="...@....iam.gserviceaccount.com",
|
||||
client_id="your_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/..."
|
||||
),
|
||||
universe_domain="googleapis.com",
|
||||
),
|
||||
scopes=["https://www.googleapis.com/auth/cloud-platform"],
|
||||
service_account_credential=sa_credential,
|
||||
scopes=_DEFAULT_SCOPES,
|
||||
),
|
||||
)
|
||||
|
||||
@@ -95,7 +105,7 @@ def test_exchange_credential_success(
|
||||
assert result.auth_type == AuthCredentialTypes.HTTP
|
||||
assert result.http.scheme == "bearer"
|
||||
assert result.http.credentials.token == "mock_access_token"
|
||||
mock_from_service_account_info.assert_called_once()
|
||||
mock_from_sa_info.assert_called_once()
|
||||
mock_credentials.refresh.assert_called_once()
|
||||
|
||||
|
||||
@@ -107,7 +117,7 @@ def test_exchange_credential_success(
|
||||
(None, None, None),
|
||||
],
|
||||
)
|
||||
def test_exchange_credential_use_default_credential_success(
|
||||
def test_exchange_access_token_with_adc_sets_quota_project(
|
||||
service_account_exchanger,
|
||||
auth_scheme,
|
||||
monkeypatch,
|
||||
@@ -115,7 +125,6 @@ def test_exchange_credential_use_default_credential_success(
|
||||
adc_project_id,
|
||||
expected_quota_project_id,
|
||||
):
|
||||
"""Test successful exchange of service account credentials using default credential."""
|
||||
mock_credentials = MagicMock()
|
||||
mock_credentials.token = "mock_access_token"
|
||||
mock_credentials.quota_project_id = cred_quota_project_id
|
||||
@@ -128,7 +137,7 @@ def test_exchange_credential_use_default_credential_success(
|
||||
auth_type=AuthCredentialTypes.SERVICE_ACCOUNT,
|
||||
service_account=ServiceAccount(
|
||||
use_default_credential=True,
|
||||
scopes=["https://www.googleapis.com/auth/cloud-platform"],
|
||||
scopes=["https://www.googleapis.com/auth/bigquery"],
|
||||
),
|
||||
)
|
||||
|
||||
@@ -146,26 +155,49 @@ def test_exchange_credential_use_default_credential_success(
|
||||
)
|
||||
else:
|
||||
assert not result.http.additional_headers
|
||||
# Verify google.auth.default is called with the correct scopes parameter
|
||||
mock_google_auth_default.assert_called_once_with(
|
||||
scopes=["https://www.googleapis.com/auth/cloud-platform"]
|
||||
scopes=["https://www.googleapis.com/auth/bigquery"]
|
||||
)
|
||||
mock_credentials.refresh.assert_called_once()
|
||||
|
||||
|
||||
def test_exchange_credential_missing_auth_credential(
|
||||
def test_exchange_access_token_with_adc_defaults_to_cloud_platform_scope(
|
||||
service_account_exchanger, auth_scheme, monkeypatch
|
||||
):
|
||||
mock_credentials = MagicMock()
|
||||
mock_credentials.token = "mock_access_token"
|
||||
mock_credentials.quota_project_id = None
|
||||
mock_google_auth_default = MagicMock(return_value=(mock_credentials, None))
|
||||
monkeypatch.setattr(google.auth, "default", mock_google_auth_default)
|
||||
|
||||
auth_credential = AuthCredential(
|
||||
auth_type=AuthCredentialTypes.SERVICE_ACCOUNT,
|
||||
service_account=ServiceAccount(
|
||||
use_default_credential=True,
|
||||
),
|
||||
)
|
||||
|
||||
result = service_account_exchanger.exchange_credential(
|
||||
auth_scheme, auth_credential
|
||||
)
|
||||
|
||||
assert result.auth_type == AuthCredentialTypes.HTTP
|
||||
assert result.http.scheme == "bearer"
|
||||
assert result.http.credentials.token == "mock_access_token"
|
||||
mock_google_auth_default.assert_called_once_with(scopes=_DEFAULT_SCOPES)
|
||||
|
||||
|
||||
def test_exchange_raises_when_auth_credential_is_none(
|
||||
service_account_exchanger, auth_scheme
|
||||
):
|
||||
"""Test missing auth credential during exchange."""
|
||||
with pytest.raises(AuthCredentialMissingError) as exc_info:
|
||||
service_account_exchanger.exchange_credential(auth_scheme, None)
|
||||
assert "Service account credentials are missing" in str(exc_info.value)
|
||||
|
||||
|
||||
def test_exchange_credential_missing_service_account_info(
|
||||
def test_exchange_raises_when_service_account_is_none(
|
||||
service_account_exchanger, auth_scheme
|
||||
):
|
||||
"""Test missing service account info during exchange."""
|
||||
auth_credential = AuthCredential(
|
||||
auth_type=AuthCredentialTypes.SERVICE_ACCOUNT,
|
||||
)
|
||||
@@ -174,47 +206,188 @@ def test_exchange_credential_missing_service_account_info(
|
||||
assert "Service account credentials are missing" in str(exc_info.value)
|
||||
|
||||
|
||||
def test_exchange_credential_exchange_failure(
|
||||
service_account_exchanger, auth_scheme, monkeypatch
|
||||
def test_exchange_wraps_google_auth_error_as_missing_error(
|
||||
service_account_exchanger, auth_scheme, sa_credential, monkeypatch
|
||||
):
|
||||
"""Test failure during service account token exchange."""
|
||||
mock_from_service_account_info = MagicMock(
|
||||
side_effect=Exception("Failed to load credentials")
|
||||
)
|
||||
target_path = (
|
||||
"google.adk.tools.openapi_tool.auth.credential_exchangers."
|
||||
"service_account_exchanger.service_account.Credentials."
|
||||
"from_service_account_info"
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
target_path,
|
||||
mock_from_service_account_info,
|
||||
mock_from_sa_info = MagicMock(
|
||||
side_effect=ValueError("Failed to load credentials")
|
||||
)
|
||||
monkeypatch.setattr(_ACCESS_TOKEN_MONKEYPATCH_TARGET, mock_from_sa_info)
|
||||
|
||||
auth_credential = AuthCredential(
|
||||
auth_type=AuthCredentialTypes.SERVICE_ACCOUNT,
|
||||
service_account=ServiceAccount(
|
||||
service_account_credential=ServiceAccountCredential(
|
||||
type_="service_account",
|
||||
project_id="your_project_id",
|
||||
private_key_id="your_private_key_id",
|
||||
private_key="-----BEGIN PRIVATE KEY-----...",
|
||||
client_email="...@....iam.gserviceaccount.com",
|
||||
client_id="your_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/..."
|
||||
),
|
||||
universe_domain="googleapis.com",
|
||||
),
|
||||
scopes=["https://www.googleapis.com/auth/cloud-platform"],
|
||||
service_account_credential=sa_credential,
|
||||
scopes=_DEFAULT_SCOPES,
|
||||
),
|
||||
)
|
||||
|
||||
with pytest.raises(AuthCredentialMissingError) as exc_info:
|
||||
service_account_exchanger.exchange_credential(auth_scheme, auth_credential)
|
||||
assert "Failed to exchange service account token" in str(exc_info.value)
|
||||
mock_from_service_account_info.assert_called_once()
|
||||
mock_from_sa_info.assert_called_once()
|
||||
|
||||
|
||||
def test_exchange_raises_when_explicit_credentials_have_no_scopes(
|
||||
service_account_exchanger, auth_scheme, sa_credential
|
||||
):
|
||||
auth_credential = AuthCredential(
|
||||
auth_type=AuthCredentialTypes.SERVICE_ACCOUNT,
|
||||
service_account=ServiceAccount(
|
||||
service_account_credential=sa_credential,
|
||||
),
|
||||
)
|
||||
|
||||
with pytest.raises(AuthCredentialMissingError) as exc_info:
|
||||
service_account_exchanger.exchange_credential(auth_scheme, auth_credential)
|
||||
assert "scopes are required" in str(exc_info.value)
|
||||
|
||||
|
||||
# --- ID token exchange tests ---
|
||||
|
||||
|
||||
def test_exchange_id_token_with_explicit_credentials(
|
||||
service_account_exchanger, auth_scheme, sa_credential, monkeypatch
|
||||
):
|
||||
mock_id_credentials = MagicMock()
|
||||
mock_id_credentials.token = "mock_id_token"
|
||||
mock_from_sa_info = MagicMock(return_value=mock_id_credentials)
|
||||
monkeypatch.setattr(_ID_TOKEN_MONKEYPATCH_TARGET, mock_from_sa_info)
|
||||
|
||||
auth_credential = AuthCredential(
|
||||
auth_type=AuthCredentialTypes.SERVICE_ACCOUNT,
|
||||
service_account=ServiceAccount(
|
||||
service_account_credential=sa_credential,
|
||||
scopes=_DEFAULT_SCOPES,
|
||||
use_id_token=True,
|
||||
audience="https://my-service.run.app",
|
||||
),
|
||||
)
|
||||
|
||||
result = service_account_exchanger.exchange_credential(
|
||||
auth_scheme, auth_credential
|
||||
)
|
||||
|
||||
assert result.auth_type == AuthCredentialTypes.HTTP
|
||||
assert result.http.scheme == "bearer"
|
||||
assert result.http.credentials.token == "mock_id_token"
|
||||
assert result.http.additional_headers is None
|
||||
mock_from_sa_info.assert_called_once()
|
||||
assert (
|
||||
mock_from_sa_info.call_args[1]["target_audience"]
|
||||
== "https://my-service.run.app"
|
||||
)
|
||||
mock_id_credentials.refresh.assert_called_once()
|
||||
|
||||
|
||||
def test_exchange_id_token_with_adc(
|
||||
service_account_exchanger, auth_scheme, monkeypatch
|
||||
):
|
||||
mock_fetch_id_token = MagicMock(return_value="mock_adc_id_token")
|
||||
monkeypatch.setattr(_FETCH_ID_TOKEN_MONKEYPATCH_TARGET, mock_fetch_id_token)
|
||||
|
||||
auth_credential = AuthCredential(
|
||||
auth_type=AuthCredentialTypes.SERVICE_ACCOUNT,
|
||||
service_account=ServiceAccount(
|
||||
use_default_credential=True,
|
||||
scopes=_DEFAULT_SCOPES,
|
||||
use_id_token=True,
|
||||
audience="https://my-service.run.app",
|
||||
),
|
||||
)
|
||||
|
||||
result = service_account_exchanger.exchange_credential(
|
||||
auth_scheme, auth_credential
|
||||
)
|
||||
|
||||
assert result.auth_type == AuthCredentialTypes.HTTP
|
||||
assert result.http.scheme == "bearer"
|
||||
assert result.http.credentials.token == "mock_adc_id_token"
|
||||
assert result.http.additional_headers is None
|
||||
mock_fetch_id_token.assert_called_once()
|
||||
assert mock_fetch_id_token.call_args[0][1] == "https://my-service.run.app"
|
||||
|
||||
|
||||
def test_id_token_requires_audience():
|
||||
with pytest.raises(
|
||||
ValueError, match="audience is required when use_id_token is True"
|
||||
):
|
||||
ServiceAccount(
|
||||
use_default_credential=True,
|
||||
use_id_token=True,
|
||||
)
|
||||
|
||||
|
||||
def test_exchange_id_token_wraps_error_with_explicit_credentials(
|
||||
service_account_exchanger, auth_scheme, sa_credential, monkeypatch
|
||||
):
|
||||
mock_from_sa_info = MagicMock(
|
||||
side_effect=ValueError("Failed to create ID token credentials")
|
||||
)
|
||||
monkeypatch.setattr(_ID_TOKEN_MONKEYPATCH_TARGET, mock_from_sa_info)
|
||||
|
||||
auth_credential = AuthCredential(
|
||||
auth_type=AuthCredentialTypes.SERVICE_ACCOUNT,
|
||||
service_account=ServiceAccount(
|
||||
service_account_credential=sa_credential,
|
||||
scopes=_DEFAULT_SCOPES,
|
||||
use_id_token=True,
|
||||
audience="https://my-service.run.app",
|
||||
),
|
||||
)
|
||||
|
||||
with pytest.raises(AuthCredentialMissingError) as exc_info:
|
||||
service_account_exchanger.exchange_credential(auth_scheme, auth_credential)
|
||||
assert "Failed to exchange service account for ID token" in str(
|
||||
exc_info.value
|
||||
)
|
||||
|
||||
|
||||
def test_exchange_id_token_wraps_error_with_adc(
|
||||
service_account_exchanger, auth_scheme, monkeypatch
|
||||
):
|
||||
mock_fetch_id_token = MagicMock(
|
||||
side_effect=google_auth_exceptions.DefaultCredentialsError(
|
||||
"Metadata service unavailable"
|
||||
)
|
||||
)
|
||||
monkeypatch.setattr(_FETCH_ID_TOKEN_MONKEYPATCH_TARGET, mock_fetch_id_token)
|
||||
|
||||
auth_credential = AuthCredential(
|
||||
auth_type=AuthCredentialTypes.SERVICE_ACCOUNT,
|
||||
service_account=ServiceAccount(
|
||||
use_default_credential=True,
|
||||
scopes=_DEFAULT_SCOPES,
|
||||
use_id_token=True,
|
||||
audience="https://my-service.run.app",
|
||||
),
|
||||
)
|
||||
|
||||
with pytest.raises(AuthCredentialMissingError) as exc_info:
|
||||
service_account_exchanger.exchange_credential(auth_scheme, auth_credential)
|
||||
assert "Failed to exchange service account for ID token" in str(
|
||||
exc_info.value
|
||||
)
|
||||
|
||||
|
||||
# --- Model validator tests ---
|
||||
|
||||
|
||||
def test_model_validator_rejects_missing_credential_without_adc():
|
||||
with pytest.raises(
|
||||
ValueError,
|
||||
match="service_account_credential is required",
|
||||
):
|
||||
ServiceAccount(
|
||||
use_default_credential=False,
|
||||
scopes=_DEFAULT_SCOPES,
|
||||
)
|
||||
|
||||
|
||||
def test_model_validator_allows_adc_without_explicit_credential():
|
||||
sa = ServiceAccount(
|
||||
use_default_credential=True,
|
||||
scopes=_DEFAULT_SCOPES,
|
||||
)
|
||||
assert sa.service_account_credential is None
|
||||
assert sa.use_default_credential is True
|
||||
|
||||
Reference in New Issue
Block a user