mirror of
https://github.com/encounter/adk-python.git
synced 2026-07-09 18:19:28 -07:00
refactor: Adapt service account credential exchanger to base credential exchanger interface
PiperOrigin-RevId: 772710438
This commit is contained in:
committed by
Copybara-Service
parent
f9fa7841df
commit
0a9625317a
@@ -15,11 +15,9 @@
|
|||||||
"""Credential exchanger module."""
|
"""Credential exchanger module."""
|
||||||
|
|
||||||
from .base_credential_exchanger import BaseCredentialExchanger
|
from .base_credential_exchanger import BaseCredentialExchanger
|
||||||
from .credential_exchanger_registry import CredentialExchangerRegistry
|
|
||||||
from .service_account_credential_exchanger import ServiceAccountCredentialExchanger
|
from .service_account_credential_exchanger import ServiceAccountCredentialExchanger
|
||||||
|
|
||||||
__all__ = [
|
__all__ = [
|
||||||
"BaseCredentialExchanger",
|
"BaseCredentialExchanger",
|
||||||
"CredentialExchangerRegistry",
|
|
||||||
"ServiceAccountCredentialExchanger",
|
"ServiceAccountCredentialExchanger",
|
||||||
]
|
]
|
||||||
|
|||||||
+41
-29
@@ -16,19 +16,22 @@
|
|||||||
|
|
||||||
from __future__ import annotations
|
from __future__ import annotations
|
||||||
|
|
||||||
|
from typing import Optional
|
||||||
|
|
||||||
import google.auth
|
import google.auth
|
||||||
from google.auth.transport.requests import Request
|
from google.auth.transport.requests import Request
|
||||||
from google.oauth2 import service_account
|
from google.oauth2 import service_account
|
||||||
|
from typing_extensions import override
|
||||||
|
|
||||||
from ..utils.feature_decorator import experimental
|
from ...utils.feature_decorator import experimental
|
||||||
from .auth_credential import AuthCredential
|
from ..auth_credential import AuthCredential
|
||||||
from .auth_credential import AuthCredentialTypes
|
from ..auth_credential import AuthCredentialTypes
|
||||||
from .auth_credential import HttpAuth
|
from ..auth_schemes import AuthScheme
|
||||||
from .auth_credential import HttpCredentials
|
from .base_credential_exchanger import BaseCredentialExchanger
|
||||||
|
|
||||||
|
|
||||||
@experimental
|
@experimental
|
||||||
class ServiceAccountCredentialExchanger:
|
class ServiceAccountCredentialExchanger(BaseCredentialExchanger):
|
||||||
"""Exchanges Google Service Account credentials for an access token.
|
"""Exchanges Google Service Account credentials for an access token.
|
||||||
|
|
||||||
Uses the default service credential if `use_default_credential = True`.
|
Uses the default service credential if `use_default_credential = True`.
|
||||||
@@ -36,44 +39,56 @@ class ServiceAccountCredentialExchanger:
|
|||||||
credential.
|
credential.
|
||||||
"""
|
"""
|
||||||
|
|
||||||
def __init__(self, credential: AuthCredential):
|
@override
|
||||||
if credential.auth_type != AuthCredentialTypes.SERVICE_ACCOUNT:
|
async def exchange(
|
||||||
raise ValueError("Credential is not a service account credential.")
|
self,
|
||||||
self._credential = credential
|
auth_credential: AuthCredential,
|
||||||
|
auth_scheme: Optional[AuthScheme] = None,
|
||||||
def exchange(self) -> AuthCredential:
|
) -> AuthCredential:
|
||||||
"""Exchanges the service account auth credential for an access token.
|
"""Exchanges the service account auth credential for an access token.
|
||||||
|
|
||||||
If the AuthCredential contains a service account credential, it will be used
|
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,
|
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.
|
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:
|
Returns:
|
||||||
An AuthCredential in HTTP Bearer format, containing the access token.
|
An AuthCredential in OAUTH2 format, containing the exchanged credential JSON.
|
||||||
|
|
||||||
Raises:
|
Raises:
|
||||||
ValueError: If service account credentials are missing or invalid.
|
ValueError: If service account credentials are missing or invalid.
|
||||||
Exception: If credential exchange or refresh fails.
|
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 (
|
if (
|
||||||
self._credential is None
|
auth_credential.service_account.service_account_credential is None
|
||||||
or self._credential.service_account is None
|
and not auth_credential.service_account.use_default_credential
|
||||||
or (
|
|
||||||
self._credential.service_account.service_account_credential is None
|
|
||||||
and not self._credential.service_account.use_default_credential
|
|
||||||
)
|
|
||||||
):
|
):
|
||||||
raise ValueError(
|
raise ValueError(
|
||||||
"Service account credentials are missing. Please provide them, or set"
|
"Service account credentials are invalid. Please set the"
|
||||||
" `use_default_credential = True` to use application default"
|
" service_account_credential field or set `use_default_credential ="
|
||||||
" credential in a hosted service like Google Cloud Run."
|
" True` to use application default credential in a hosted service"
|
||||||
|
" like Google Cloud Run."
|
||||||
)
|
)
|
||||||
|
|
||||||
try:
|
try:
|
||||||
if self._credential.service_account.use_default_credential:
|
if auth_credential.service_account.use_default_credential:
|
||||||
credentials, _ = google.auth.default()
|
credentials, _ = google.auth.default()
|
||||||
else:
|
else:
|
||||||
config = self._credential.service_account
|
config = auth_credential.service_account
|
||||||
credentials = service_account.Credentials.from_service_account_info(
|
credentials = service_account.Credentials.from_service_account_info(
|
||||||
config.service_account_credential.model_dump(), scopes=config.scopes
|
config.service_account_credential.model_dump(), scopes=config.scopes
|
||||||
)
|
)
|
||||||
@@ -82,11 +97,8 @@ class ServiceAccountCredentialExchanger:
|
|||||||
credentials.refresh(Request())
|
credentials.refresh(Request())
|
||||||
|
|
||||||
return AuthCredential(
|
return AuthCredential(
|
||||||
auth_type=AuthCredentialTypes.HTTP,
|
auth_type=AuthCredentialTypes.OAUTH2,
|
||||||
http=HttpAuth(
|
google_oauth2_json=credentials.to_json(),
|
||||||
scheme="bearer",
|
|
||||||
credentials=HttpCredentials(token=credentials.token),
|
|
||||||
),
|
|
||||||
)
|
)
|
||||||
except Exception as e:
|
except Exception as e:
|
||||||
raise ValueError(f"Failed to exchange service account token: {e}") from e
|
raise ValueError(f"Failed to exchange service account token: {e}") from e
|
||||||
@@ -0,0 +1,15 @@
|
|||||||
|
# Copyright 2025 Google LLC
|
||||||
|
#
|
||||||
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||||
|
# you may not use this file except in compliance with the License.
|
||||||
|
# You may obtain a copy of the License at
|
||||||
|
#
|
||||||
|
# http://www.apache.org/licenses/LICENSE-2.0
|
||||||
|
#
|
||||||
|
# Unless required by applicable law or agreed to in writing, software
|
||||||
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||||
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||||
|
# See the License for the specific language governing permissions and
|
||||||
|
# limitations under the License.
|
||||||
|
|
||||||
|
"""Tests for credential exchanger."""
|
||||||
+147
-55
@@ -17,19 +17,20 @@
|
|||||||
from unittest.mock import MagicMock
|
from unittest.mock import MagicMock
|
||||||
from unittest.mock import patch
|
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 AuthCredential
|
||||||
from google.adk.auth.auth_credential import AuthCredentialTypes
|
from google.adk.auth.auth_credential import AuthCredentialTypes
|
||||||
from google.adk.auth.auth_credential import ServiceAccount
|
from google.adk.auth.auth_credential import ServiceAccount
|
||||||
from google.adk.auth.auth_credential import ServiceAccountCredential
|
from google.adk.auth.auth_credential import ServiceAccountCredential
|
||||||
from google.adk.auth.service_account_credential_exchanger import ServiceAccountCredentialExchanger
|
from google.adk.auth.exchanger.service_account_credential_exchanger import ServiceAccountCredentialExchanger
|
||||||
import pytest
|
import pytest
|
||||||
|
|
||||||
|
|
||||||
class TestServiceAccountCredentialExchanger:
|
class TestServiceAccountCredentialExchanger:
|
||||||
"""Test cases for ServiceAccountCredentialExchanger."""
|
"""Test cases for ServiceAccountCredentialExchanger."""
|
||||||
|
|
||||||
def test_init_valid_credential(self):
|
def test_exchange_with_valid_credential(self):
|
||||||
"""Test successful initialization with valid service account credential."""
|
"""Test successful exchange with valid service account credential."""
|
||||||
credential = AuthCredential(
|
credential = AuthCredential(
|
||||||
auth_type=AuthCredentialTypes.SERVICE_ACCOUNT,
|
auth_type=AuthCredentialTypes.SERVICE_ACCOUNT,
|
||||||
service_account=ServiceAccount(
|
service_account=ServiceAccount(
|
||||||
@@ -55,26 +56,36 @@ class TestServiceAccountCredentialExchanger:
|
|||||||
),
|
),
|
||||||
)
|
)
|
||||||
|
|
||||||
exchanger = ServiceAccountCredentialExchanger(credential)
|
auth_scheme = HTTPBearer()
|
||||||
assert exchanger._credential == credential
|
exchanger = ServiceAccountCredentialExchanger()
|
||||||
|
|
||||||
def test_init_invalid_credential_type(self):
|
# This should not raise an exception
|
||||||
"""Test initialization with invalid credential type raises ValueError."""
|
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(
|
credential = AuthCredential(
|
||||||
auth_type=AuthCredentialTypes.API_KEY,
|
auth_type=AuthCredentialTypes.API_KEY,
|
||||||
api_key="test-key",
|
api_key="test-key",
|
||||||
)
|
)
|
||||||
|
|
||||||
|
auth_scheme = HTTPBearer()
|
||||||
|
exchanger = ServiceAccountCredentialExchanger()
|
||||||
|
|
||||||
with pytest.raises(
|
with pytest.raises(
|
||||||
ValueError, match="Credential is not a service account credential"
|
ValueError, match="Credential is not a service account credential"
|
||||||
):
|
):
|
||||||
ServiceAccountCredentialExchanger(credential)
|
await exchanger.exchange(credential, auth_scheme)
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
@patch(
|
@patch(
|
||||||
"google.adk.auth.service_account_credential_exchanger.service_account.Credentials.from_service_account_info"
|
"google.adk.auth.exchanger.service_account_credential_exchanger.service_account.Credentials.from_service_account_info"
|
||||||
)
|
)
|
||||||
@patch("google.adk.auth.service_account_credential_exchanger.Request")
|
@patch(
|
||||||
def test_exchange_with_explicit_credentials_success(
|
"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
|
self, mock_request_class, mock_from_service_account_info
|
||||||
):
|
):
|
||||||
"""Test successful exchange with explicit service account credentials."""
|
"""Test successful exchange with explicit service account credentials."""
|
||||||
@@ -84,6 +95,9 @@ class TestServiceAccountCredentialExchanger:
|
|||||||
|
|
||||||
mock_credentials = MagicMock()
|
mock_credentials = MagicMock()
|
||||||
mock_credentials.token = "mock_access_token"
|
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
|
mock_from_service_account_info.return_value = mock_credentials
|
||||||
|
|
||||||
# Create test credential
|
# Create test credential
|
||||||
@@ -113,13 +127,20 @@ class TestServiceAccountCredentialExchanger:
|
|||||||
),
|
),
|
||||||
)
|
)
|
||||||
|
|
||||||
exchanger = ServiceAccountCredentialExchanger(credential)
|
auth_scheme = HTTPBearer()
|
||||||
result = exchanger.exchange()
|
exchanger = ServiceAccountCredentialExchanger()
|
||||||
|
result = await exchanger.exchange(credential, auth_scheme)
|
||||||
|
|
||||||
# Verify the result
|
# Verify the result
|
||||||
assert result.auth_type == AuthCredentialTypes.HTTP
|
assert result.auth_type == AuthCredentialTypes.OAUTH2
|
||||||
assert result.http.scheme == "bearer"
|
assert result.google_oauth2_json is not None
|
||||||
assert result.http.credentials.token == "mock_access_token"
|
# 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
|
# Verify mocks were called correctly
|
||||||
mock_from_service_account_info.assert_called_once_with(
|
mock_from_service_account_info.assert_called_once_with(
|
||||||
@@ -128,11 +149,14 @@ class TestServiceAccountCredentialExchanger:
|
|||||||
)
|
)
|
||||||
mock_credentials.refresh.assert_called_once_with(mock_request)
|
mock_credentials.refresh.assert_called_once_with(mock_request)
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
@patch(
|
@patch(
|
||||||
"google.adk.auth.service_account_credential_exchanger.google.auth.default"
|
"google.adk.auth.exchanger.service_account_credential_exchanger.google.auth.default"
|
||||||
)
|
)
|
||||||
@patch("google.adk.auth.service_account_credential_exchanger.Request")
|
@patch(
|
||||||
def test_exchange_with_default_credentials_success(
|
"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
|
self, mock_request_class, mock_google_auth_default
|
||||||
):
|
):
|
||||||
"""Test successful exchange with default application credentials."""
|
"""Test successful exchange with default application credentials."""
|
||||||
@@ -142,6 +166,9 @@ class TestServiceAccountCredentialExchanger:
|
|||||||
|
|
||||||
mock_credentials = MagicMock()
|
mock_credentials = MagicMock()
|
||||||
mock_credentials.token = "default_access_token"
|
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")
|
mock_google_auth_default.return_value = (mock_credentials, "test-project")
|
||||||
|
|
||||||
# Create test credential with use_default_credential=True
|
# Create test credential with use_default_credential=True
|
||||||
@@ -153,33 +180,45 @@ class TestServiceAccountCredentialExchanger:
|
|||||||
),
|
),
|
||||||
)
|
)
|
||||||
|
|
||||||
exchanger = ServiceAccountCredentialExchanger(credential)
|
auth_scheme = HTTPBearer()
|
||||||
result = exchanger.exchange()
|
exchanger = ServiceAccountCredentialExchanger()
|
||||||
|
result = await exchanger.exchange(credential, auth_scheme)
|
||||||
|
|
||||||
# Verify the result
|
# Verify the result
|
||||||
assert result.auth_type == AuthCredentialTypes.HTTP
|
assert result.auth_type == AuthCredentialTypes.OAUTH2
|
||||||
assert result.http.scheme == "bearer"
|
assert result.google_oauth2_json is not None
|
||||||
assert result.http.credentials.token == "default_access_token"
|
# 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
|
# Verify mocks were called correctly
|
||||||
mock_google_auth_default.assert_called_once()
|
mock_google_auth_default.assert_called_once()
|
||||||
mock_credentials.refresh.assert_called_once_with(mock_request)
|
mock_credentials.refresh.assert_called_once_with(mock_request)
|
||||||
|
|
||||||
def test_exchange_missing_service_account(self):
|
@pytest.mark.asyncio
|
||||||
|
async def test_exchange_missing_service_account(self):
|
||||||
"""Test exchange fails when service_account is None."""
|
"""Test exchange fails when service_account is None."""
|
||||||
credential = AuthCredential(
|
credential = AuthCredential(
|
||||||
auth_type=AuthCredentialTypes.SERVICE_ACCOUNT,
|
auth_type=AuthCredentialTypes.SERVICE_ACCOUNT,
|
||||||
service_account=None,
|
service_account=None,
|
||||||
)
|
)
|
||||||
|
|
||||||
exchanger = ServiceAccountCredentialExchanger(credential)
|
auth_scheme = HTTPBearer()
|
||||||
|
exchanger = ServiceAccountCredentialExchanger()
|
||||||
|
|
||||||
with pytest.raises(
|
with pytest.raises(
|
||||||
ValueError, match="Service account credentials are missing"
|
ValueError, match="Service account credentials are missing"
|
||||||
):
|
):
|
||||||
exchanger.exchange()
|
await exchanger.exchange(credential, auth_scheme)
|
||||||
|
|
||||||
def test_exchange_missing_credentials_and_not_default(self):
|
@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."""
|
"""Test exchange fails when credentials are missing and use_default_credential is False."""
|
||||||
credential = AuthCredential(
|
credential = AuthCredential(
|
||||||
auth_type=AuthCredentialTypes.SERVICE_ACCOUNT,
|
auth_type=AuthCredentialTypes.SERVICE_ACCOUNT,
|
||||||
@@ -190,17 +229,19 @@ class TestServiceAccountCredentialExchanger:
|
|||||||
),
|
),
|
||||||
)
|
)
|
||||||
|
|
||||||
exchanger = ServiceAccountCredentialExchanger(credential)
|
auth_scheme = HTTPBearer()
|
||||||
|
exchanger = ServiceAccountCredentialExchanger()
|
||||||
|
|
||||||
with pytest.raises(
|
with pytest.raises(
|
||||||
ValueError, match="Service account credentials are missing"
|
ValueError, match="Service account credentials are invalid"
|
||||||
):
|
):
|
||||||
exchanger.exchange()
|
await exchanger.exchange(credential, auth_scheme)
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
@patch(
|
@patch(
|
||||||
"google.adk.auth.service_account_credential_exchanger.service_account.Credentials.from_service_account_info"
|
"google.adk.auth.exchanger.service_account_credential_exchanger.service_account.Credentials.from_service_account_info"
|
||||||
)
|
)
|
||||||
def test_exchange_credential_creation_failure(
|
async def test_exchange_credential_creation_failure(
|
||||||
self, mock_from_service_account_info
|
self, mock_from_service_account_info
|
||||||
):
|
):
|
||||||
"""Test exchange handles credential creation failure gracefully."""
|
"""Test exchange handles credential creation failure gracefully."""
|
||||||
@@ -234,17 +275,21 @@ class TestServiceAccountCredentialExchanger:
|
|||||||
),
|
),
|
||||||
)
|
)
|
||||||
|
|
||||||
exchanger = ServiceAccountCredentialExchanger(credential)
|
auth_scheme = HTTPBearer()
|
||||||
|
exchanger = ServiceAccountCredentialExchanger()
|
||||||
|
|
||||||
with pytest.raises(
|
with pytest.raises(
|
||||||
ValueError, match="Failed to exchange service account token"
|
ValueError, match="Failed to exchange service account token"
|
||||||
):
|
):
|
||||||
exchanger.exchange()
|
await exchanger.exchange(credential, auth_scheme)
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
@patch(
|
@patch(
|
||||||
"google.adk.auth.service_account_credential_exchanger.google.auth.default"
|
"google.adk.auth.exchanger.service_account_credential_exchanger.google.auth.default"
|
||||||
)
|
)
|
||||||
def test_exchange_default_credential_failure(self, mock_google_auth_default):
|
async def test_exchange_default_credential_failure(
|
||||||
|
self, mock_google_auth_default
|
||||||
|
):
|
||||||
"""Test exchange handles default credential failure gracefully."""
|
"""Test exchange handles default credential failure gracefully."""
|
||||||
# Setup mock to raise exception
|
# Setup mock to raise exception
|
||||||
mock_google_auth_default.side_effect = Exception(
|
mock_google_auth_default.side_effect = Exception(
|
||||||
@@ -260,18 +305,22 @@ class TestServiceAccountCredentialExchanger:
|
|||||||
),
|
),
|
||||||
)
|
)
|
||||||
|
|
||||||
exchanger = ServiceAccountCredentialExchanger(credential)
|
auth_scheme = HTTPBearer()
|
||||||
|
exchanger = ServiceAccountCredentialExchanger()
|
||||||
|
|
||||||
with pytest.raises(
|
with pytest.raises(
|
||||||
ValueError, match="Failed to exchange service account token"
|
ValueError, match="Failed to exchange service account token"
|
||||||
):
|
):
|
||||||
exchanger.exchange()
|
await exchanger.exchange(credential, auth_scheme)
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
@patch(
|
@patch(
|
||||||
"google.adk.auth.service_account_credential_exchanger.service_account.Credentials.from_service_account_info"
|
"google.adk.auth.exchanger.service_account_credential_exchanger.service_account.Credentials.from_service_account_info"
|
||||||
)
|
)
|
||||||
@patch("google.adk.auth.service_account_credential_exchanger.Request")
|
@patch(
|
||||||
def test_exchange_refresh_failure(
|
"google.adk.auth.exchanger.service_account_credential_exchanger.Request"
|
||||||
|
)
|
||||||
|
async def test_exchange_refresh_failure(
|
||||||
self, mock_request_class, mock_from_service_account_info
|
self, mock_request_class, mock_from_service_account_info
|
||||||
):
|
):
|
||||||
"""Test exchange handles credential refresh failure gracefully."""
|
"""Test exchange handles credential refresh failure gracefully."""
|
||||||
@@ -312,30 +361,73 @@ class TestServiceAccountCredentialExchanger:
|
|||||||
),
|
),
|
||||||
)
|
)
|
||||||
|
|
||||||
exchanger = ServiceAccountCredentialExchanger(credential)
|
auth_scheme = HTTPBearer()
|
||||||
|
exchanger = ServiceAccountCredentialExchanger()
|
||||||
|
|
||||||
with pytest.raises(
|
with pytest.raises(
|
||||||
ValueError, match="Failed to exchange service account token"
|
ValueError, match="Failed to exchange service account token"
|
||||||
):
|
):
|
||||||
exchanger.exchange()
|
await exchanger.exchange(credential, auth_scheme)
|
||||||
|
|
||||||
def test_exchange_none_credential_in_constructor(self):
|
@pytest.mark.asyncio
|
||||||
"""Test that passing None credential raises appropriate error during construction."""
|
async def test_exchange_none_credential_in_constructor(self):
|
||||||
# This test verifies behavior when _credential is None, though this shouldn't
|
"""Test that passing None credential raises appropriate error during exchange."""
|
||||||
# happen in normal usage due to constructor validation
|
# 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(
|
credential = AuthCredential(
|
||||||
auth_type=AuthCredentialTypes.SERVICE_ACCOUNT,
|
auth_type=AuthCredentialTypes.SERVICE_ACCOUNT,
|
||||||
service_account=ServiceAccount(
|
service_account=ServiceAccount(
|
||||||
|
service_account_credential=None,
|
||||||
use_default_credential=True,
|
use_default_credential=True,
|
||||||
scopes=["https://www.googleapis.com/auth/cloud-platform"],
|
scopes=["https://www.googleapis.com/auth/cloud-platform"],
|
||||||
),
|
),
|
||||||
)
|
)
|
||||||
|
|
||||||
exchanger = ServiceAccountCredentialExchanger(credential)
|
auth_scheme = HTTPBearer()
|
||||||
# Manually set to None to test the validation logic
|
exchanger = ServiceAccountCredentialExchanger()
|
||||||
exchanger._credential = None
|
result = await exchanger.exchange(credential, auth_scheme)
|
||||||
|
|
||||||
with pytest.raises(
|
# Verify the result
|
||||||
ValueError, match="Service account credentials are missing"
|
assert result.auth_type == AuthCredentialTypes.OAUTH2
|
||||||
):
|
assert result.google_oauth2_json is not None
|
||||||
exchanger.exchange()
|
# 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)
|
||||||
Reference in New Issue
Block a user