chore: import upstream snapshot with attribution
Continuous Integration / Pre-commit Linter (push) Has been cancelled
Continuous Integration / Mypy Check (Python 3.10) (push) Has been cancelled
Continuous Integration / Mypy Check (Python 3.11) (push) Has been cancelled
Continuous Integration / Mypy Check (Python 3.12) (push) Has been cancelled
Continuous Integration / Mypy Check (Python 3.13) (push) Has been cancelled
Continuous Integration / Unit Tests (Python 3.10) (push) Has been cancelled
Continuous Integration / Unit Tests (Python 3.11) (push) Has been cancelled
Continuous Integration / Unit Tests (Python 3.12) (push) Has been cancelled
Continuous Integration / Unit Tests (Python 3.13) (push) Has been cancelled
Continuous Integration / Unit Tests (Python 3.14) (push) Has been cancelled
Continuous Integration / A2A v0.3 Tests (Python 3.10) (push) Has been cancelled
Continuous Integration / A2A v0.3 Tests (Python 3.11) (push) Has been cancelled
Continuous Integration / A2A v0.3 Tests (Python 3.12) (push) Has been cancelled
Copybara PR Handler / close-imported-pr (push) Has been cancelled
Continuous Integration / A2A v0.3 Tests (Python 3.13) (push) Has been cancelled
Continuous Integration / A2A v0.3 Tests (Python 3.14) (push) Has been cancelled

This commit is contained in:
wehub-resource-sync
2026-07-13 13:25:13 +08:00
commit ec2b666284
2231 changed files with 491535 additions and 0 deletions
@@ -0,0 +1,423 @@
# Copyright 2026 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.
from unittest.mock import Mock
from unittest.mock import patch
import pytest
pytest.importorskip(
"google.cloud.agentidentitycredentials_v1",
reason="Requires google-cloud-agentidentitycredentials",
)
from google.adk.agents.callback_context import CallbackContext
from google.adk.auth.auth_credential import AuthCredentialTypes
from google.adk.auth.auth_tool import AuthConfig
from google.adk.auth.auth_tool import AuthToolArguments
from google.adk.flows.llm_flows.functions import REQUEST_EUC_FUNCTION_CALL_NAME
from google.adk.integrations.agent_identity import _agent_identity_credentials_provider
from google.adk.integrations.agent_identity import GcpAuthProviderScheme
from google.adk.integrations.agent_identity._agent_identity_credentials_provider import _AgentIdentityCredentialsProvider
from google.adk.integrations.agent_identity._agent_identity_credentials_provider import Client
from google.adk.sessions.session import Session
from google.cloud.agentidentitycredentials_v1 import RetrieveCredentialsResponse
@pytest.fixture
def mock_client():
return Mock(spec=Client)
@pytest.fixture
def provider(mock_client):
return _AgentIdentityCredentialsProvider(client=mock_client)
@pytest.fixture
def auth_scheme():
scheme = GcpAuthProviderScheme(
name="projects/test-project/locations/global/authProviders/test-provider",
scopes=["test-scope"],
continue_uri="https://example.com/continue",
)
return scheme
@pytest.fixture
def mock_response(mock_client):
resp = RetrieveCredentialsResponse()
mock_client.retrieve_credentials.return_value = resp
return resp
@pytest.fixture
def context():
context = Mock(spec=CallbackContext)
context.user_id = "user"
context.function_call_id = "call_123"
session = Mock(spec=Session)
session.events = []
context.session = session
return context
@patch.dict(_agent_identity_credentials_provider.os.environ, clear=True)
@patch.object(_agent_identity_credentials_provider, "Client")
def test_get_client_uses_rest_transport(mock_client_class):
provider = (
_agent_identity_credentials_provider._AgentIdentityCredentialsProvider()
)
provider._get_client()
mock_client_class.assert_called_once()
_, kwargs = mock_client_class.call_args
assert kwargs.get("transport") == "rest"
@patch.dict(
_agent_identity_credentials_provider.os.environ,
{"AGENT_IDENTITY_CREDENTIALS_TARGET_HOST": "some-host"},
)
@patch.object(_agent_identity_credentials_provider, "Client")
@patch.object(_agent_identity_credentials_provider, "ClientOptions")
def test_get_client_with_env_var(mock_client_options_class, mock_client_class):
provider = (
_agent_identity_credentials_provider._AgentIdentityCredentialsProvider()
)
client = provider._get_client()
assert client == mock_client_class.return_value
mock_client_options_class.assert_called_once_with(api_endpoint="some-host")
mock_client_class.assert_called_once_with(
client_options=mock_client_options_class.return_value, transport="rest"
)
# ==============================================================================
# Non-interactive auth flows (API key and 2-legged OAuth)
# ==============================================================================
async def test_get_auth_credential_raises_error_if_context_is_missing(
provider, auth_scheme
):
"""Test get_auth_credential raises ValueError if context is missing."""
with pytest.raises(
ValueError,
match="GcpAuthProvider requires a context with a valid user_id",
):
await provider.get_auth_credential(auth_scheme, context=None)
async def test_get_auth_credential_raises_error_if_user_id_is_missing(
provider, auth_scheme
):
"""Test get_auth_credential raises ValueError if user_id is missing."""
context = Mock(spec=CallbackContext)
context.user_id = None
with pytest.raises(
ValueError,
match="GcpAuthProvider requires a context with a valid user_id",
):
await provider.get_auth_credential(auth_scheme, context=context)
async def test_get_auth_credential_returns_credential_if_available_immediately(
mock_client,
auth_scheme,
context,
provider,
):
"""Test get_auth_credential returns credential if available immediately."""
mock_response = RetrieveCredentialsResponse(
{"success": {"header": "Authorization: Bearer", "token": "test-token"}}
)
mock_client.retrieve_credentials.return_value = mock_response
auth_credential = await provider.get_auth_credential(auth_scheme, context)
assert auth_credential.auth_type == AuthCredentialTypes.HTTP
assert auth_credential.http.scheme == "Bearer"
assert auth_credential.http.credentials.token == "test-token"
mock_client.retrieve_credentials.assert_called_once()
async def test_get_auth_credential_raises_error_if_upstream_returns_empty_header(
mock_client,
auth_scheme,
context,
provider,
):
"""Test get_auth_credential raises ValueError for empty header."""
mock_response = RetrieveCredentialsResponse(
{"success": {"header": "", "token": "test-token"}}
)
mock_client.retrieve_credentials.return_value = mock_response
with pytest.raises(
ValueError,
match=(
"Received either empty header or token from Agent Identity"
" Credentials service."
),
):
await provider.get_auth_credential(auth_scheme, context)
async def test_get_auth_credential_raises_error_if_upstream_returns_empty_token(
mock_client,
auth_scheme,
context,
provider,
):
"""Test get_auth_credential raises ValueError for empty token."""
mock_response = RetrieveCredentialsResponse(
{"success": {"header": "Authorization: Bearer", "token": ""}}
)
mock_client.retrieve_credentials.return_value = mock_response
with pytest.raises(
ValueError,
match=(
"Received either empty header or token from Agent Identity"
" Credentials service."
),
):
await provider.get_auth_credential(auth_scheme, context)
async def test_get_auth_credential_returns_credential_if_upstream_returns_custom_header(
mock_client,
auth_scheme,
context,
provider,
):
"""Test get_auth_credential returns valid credential for custom header and sets X-GOOG-API-KEY header."""
mock_response = RetrieveCredentialsResponse(
{"success": {"header": "some-x-api-key", "token": "test-token"}}
)
mock_client.retrieve_credentials.return_value = mock_response
auth_credential = await provider.get_auth_credential(auth_scheme, context)
assert auth_credential.auth_type == AuthCredentialTypes.HTTP
assert not auth_credential.http.scheme
assert auth_credential.http.credentials.token is None
assert auth_credential.http.additional_headers == {
"some-x-api-key": "test-token",
"X-GOOG-API-KEY": "test-token",
}
async def test_get_auth_credential_raises_error_if_upstream_operation_errors(
mock_client, auth_scheme, context, provider
):
"""Test get_auth_credential raises RuntimeError for rejected operations."""
mock_client.retrieve_credentials.return_value = RetrieveCredentialsResponse(
{"consent_rejected": {}}
)
with pytest.raises(
RuntimeError, match="Operation failed: User consent rejected."
):
await provider.get_auth_credential(auth_scheme, context)
async def test_get_auth_credential_raises_error_if_upstream_call_fails(
mock_client, auth_scheme, context, provider
):
"""Test get_auth_credential raises RuntimeError for failed calls."""
mock_client.retrieve_credentials.side_effect = Exception(
"API Quota Exhausted"
)
with pytest.raises(
RuntimeError,
match="Failed to retrieve credential for user 'user' on provider",
) as exc_info:
await provider.get_auth_credential(auth_scheme, context)
# Assert that the original Exception is the chained cause!
assert str(exc_info.value.__cause__) == "API Quota Exhausted"
@patch.object(_agent_identity_credentials_provider.time, "time")
async def test_get_auth_credential_raises_error_if_polling_times_out(
mock_time,
mock_client,
auth_scheme,
context,
provider,
):
"""Test get_auth_credential raises RuntimeError if polling times out."""
# First call sets start_time=0.0, second call checks time > timeout
mock_time.side_effect = [0.0, 20.0]
mock_client.retrieve_credentials.return_value = RetrieveCredentialsResponse(
{"pending": {}}
)
with pytest.raises(
RuntimeError,
match="Failed to retrieve credential for user 'user' on provider",
) as exc_info:
await provider.get_auth_credential(auth_scheme, context)
assert "Timeout waiting for credentials." in str(exc_info.value.__cause__)
# ==============================================================================
# Interactive Auth Flows (3-legged OAuth for User Consents)
# ==============================================================================
async def test_get_auth_credential_initiates_user_consent(
mock_client, auth_scheme, context, provider
):
expected_uri = "https://example.com/auth"
expected_nonce = "sample-nonce-123"
mock_client.retrieve_credentials.return_value = RetrieveCredentialsResponse({
"uri_consent_required": {
"authorization_uri": expected_uri,
"consent_nonce": expected_nonce,
}
})
# Assert that there is no prior user consent completion event
assert not context.session.events
credential = await provider.get_auth_credential(auth_scheme, context)
assert credential is not None
assert credential.auth_type == AuthCredentialTypes.OAUTH2
assert credential.oauth2.auth_uri == expected_uri
assert credential.oauth2.nonce == expected_nonce
async def test_get_auth_credential_returns_fresh_auth_uri_for_repeated_requests(
mock_client, auth_scheme, context, provider
):
"""Test that repeated calls fetch fresh auth URIs if consent is still pending."""
# Arrange: Explicit initial URI
initial_uri = "https://example.com/auth"
initial_nonce = "initial-nonce-123"
mock_client.retrieve_credentials.return_value = RetrieveCredentialsResponse({
"uri_consent_required": {
"authorization_uri": initial_uri,
"consent_nonce": initial_nonce,
}
})
credential1 = await provider.get_auth_credential(auth_scheme, context)
assert credential1.oauth2.auth_uri == initial_uri
assert credential1.oauth2.nonce == initial_nonce
# Arrange: Explicit new URI for the second call
fresh_auth_uri = "https://example.com/auth_new"
fresh_nonce = "fresh-nonce-456"
mock_client.retrieve_credentials.return_value = RetrieveCredentialsResponse({
"uri_consent_required": {
"authorization_uri": fresh_auth_uri,
"consent_nonce": fresh_nonce,
}
})
credential2 = await provider.get_auth_credential(auth_scheme, context)
assert mock_client.retrieve_credentials.call_count == 2
assert credential2.oauth2.auth_uri == fresh_auth_uri
assert credential2.oauth2.nonce == fresh_nonce
async def test_get_auth_credential_returns_token_if_consent_was_completed(
mock_client, auth_scheme, context, provider
):
mock_client.retrieve_credentials.return_value = RetrieveCredentialsResponse(
{"success": {"header": "Authorization: Bearer", "token": "test-token"}}
)
# Create mock events
function_call = Mock()
function_call.id = "auth-req-1"
function_call.name = REQUEST_EUC_FUNCTION_CALL_NAME
function_call.args = AuthToolArguments(
function_call_id="call-123",
auth_config=Mock(spec=AuthConfig, auth_scheme=auth_scheme),
).model_dump(by_alias=True, exclude_none=True)
event1 = Mock()
event1.get_function_calls.return_value = [function_call]
event1.get_function_responses.return_value = []
function_response = Mock()
function_response.id = "auth-req-1"
function_response.name = REQUEST_EUC_FUNCTION_CALL_NAME
event2 = Mock()
event2.get_function_calls.return_value = []
event2.get_function_responses.return_value = [function_response]
# Setup tool context and event history (order of events matters)
context.session.events = [event1, event2]
context.function_call_id = "call-123"
# Execute
auth_credential = await provider.get_auth_credential(auth_scheme, context)
# Verify
assert auth_credential is not None
assert auth_credential.auth_type == AuthCredentialTypes.HTTP
assert auth_credential.http.scheme == "Bearer"
assert auth_credential.http.credentials.token == "test-token"
async def test_get_auth_credential_raises_error_if_consent_canceled(
mock_client, auth_scheme, context, provider
):
function_call = Mock()
function_call.id = "auth-req-1"
function_call.name = REQUEST_EUC_FUNCTION_CALL_NAME
function_call.args = AuthToolArguments(
function_call_id="call-123",
auth_config=Mock(spec=AuthConfig, auth_scheme=auth_scheme),
).model_dump(by_alias=True, exclude_none=True)
event1 = Mock()
event1.get_function_calls.return_value = [function_call]
event1.get_function_responses.return_value = []
function_response = Mock()
function_response.id = "auth-req-1"
function_response.name = REQUEST_EUC_FUNCTION_CALL_NAME
event2 = Mock()
event2.get_function_calls.return_value = []
event2.get_function_responses.return_value = [function_response]
context.session.events = [event1, event2]
context.function_call_id = "call-123"
mock_client.retrieve_credentials.return_value = RetrieveCredentialsResponse({
"uri_consent_required": {
"authorization_uri": "https://example.com/auth",
"consent_nonce": "sample-nonce",
}
})
with pytest.raises(
RuntimeError, match="Failed to retrieve consent based credential."
):
await provider.get_auth_credential(auth_scheme, context)
@@ -0,0 +1,111 @@
# Copyright 2026 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 GcpAuthProvider class."""
from unittest.mock import AsyncMock
from unittest.mock import Mock
from unittest.mock import patch
from google.adk.agents.callback_context import CallbackContext
from google.adk.auth.auth_credential import AuthCredential
from google.adk.auth.auth_tool import AuthConfig
from google.adk.integrations.agent_identity import GcpAuthProvider
from google.adk.integrations.agent_identity import GcpAuthProviderScheme
import pytest
@pytest.fixture
def auth_config():
config = Mock(spec=AuthConfig)
config.auth_scheme = Mock(spec=GcpAuthProviderScheme)
return config
@pytest.fixture
def context():
context = Mock(spec=CallbackContext)
context.user_id = "user"
return context
@pytest.fixture
def gcp_auth_provider():
return GcpAuthProvider()
def test_supported_auth_schemes(gcp_auth_provider):
"""Verify the provider supports the correct auth scheme."""
assert GcpAuthProviderScheme in gcp_auth_provider.supported_auth_schemes
async def test_get_auth_credential_raises_error_for_invalid_auth_scheme(
context,
):
"""Test get_auth_credential raises ValueError for invalid auth scheme."""
provider = GcpAuthProvider()
invalid_auth_config = Mock(spec=AuthConfig)
invalid_auth_config.auth_scheme = Mock() # Not GcpAuthProviderScheme
with pytest.raises(ValueError, match="Expected GcpAuthProviderScheme, got"):
await provider.get_auth_credential(invalid_auth_config, context)
@patch(
"google.adk.integrations.agent_identity.gcp_auth_provider._IamConnectorCredentialsProvider"
)
async def test_get_auth_credential_routes_to_iam_connector_service_provider(
mock_iam_cls, auth_config, context
):
"""Test routing to IAM Connector Credentials service for legacy auth provider resource names."""
auth_config.auth_scheme.name = (
"projects/test-project/locations/test-location/connectors/test-connector"
)
provider = GcpAuthProvider()
mock_credential = Mock(spec=AuthCredential)
mock_iam_provider = mock_iam_cls.return_value
mock_iam_provider.get_auth_credential = AsyncMock(
return_value=mock_credential
)
result = await provider.get_auth_credential(auth_config, context)
assert result == mock_credential
mock_iam_provider.get_auth_credential.assert_awaited_once_with(
auth_scheme=auth_config.auth_scheme, context=context
)
@patch(
"google.adk.integrations.agent_identity.gcp_auth_provider._AgentIdentityCredentialsProvider"
)
async def test_get_auth_credential_routes_to_agent_identity_service_provider(
mock_agent_cls, auth_config, context
):
"""Test routing to Agent Identity Credentials service for new auth provider resource names."""
auth_config.auth_scheme.name = "projects/test-project/locations/test-location/authProviders/test-provider"
provider = GcpAuthProvider()
mock_credential = Mock(spec=AuthCredential)
mock_agent_provider = mock_agent_cls.return_value
mock_agent_provider.get_auth_credential = AsyncMock(
return_value=mock_credential
)
result = await provider.get_auth_credential(auth_config, context)
assert result == mock_credential
mock_agent_provider.get_auth_credential.assert_awaited_once_with(
auth_scheme=auth_config.auth_scheme, context=context
)
@@ -0,0 +1,510 @@
# Copyright 2026 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.
from unittest.mock import Mock
from unittest.mock import patch
import pytest
pytest.importorskip(
"google.cloud.iamconnectorcredentials_v1alpha",
reason="Requires google-cloud-iamconnectorcredentials",
)
from google.adk.agents.callback_context import CallbackContext
from google.adk.auth.auth_credential import AuthCredentialTypes
from google.adk.auth.auth_tool import AuthConfig
from google.adk.auth.auth_tool import AuthToolArguments
from google.adk.flows.llm_flows.functions import REQUEST_EUC_FUNCTION_CALL_NAME
from google.adk.integrations.agent_identity import _iam_connector_credentials_provider
from google.adk.integrations.agent_identity import GcpAuthProviderScheme
from google.adk.integrations.agent_identity._iam_connector_credentials_provider import _IamConnectorCredentialsProvider
from google.adk.integrations.agent_identity._iam_connector_credentials_provider import Client
from google.adk.sessions.session import Session
from google.cloud.iamconnectorcredentials_v1alpha import RetrieveCredentialsMetadata
from google.cloud.iamconnectorcredentials_v1alpha import RetrieveCredentialsResponse
from google.longrunning.operations_pb2 import Operation
@pytest.fixture
def mock_client():
return Mock(spec=Client)
@pytest.fixture
def provider(mock_client):
return _IamConnectorCredentialsProvider(client=mock_client)
@pytest.fixture
def auth_scheme():
scheme = GcpAuthProviderScheme(
name="projects/test-project/locations/global/connectors/test-connector",
scopes=["test-scope"],
continue_uri="https://example.com/continue",
)
return scheme
@pytest.fixture
def mock_operation(mock_client):
op = Operation(done=True)
class DummyCall:
def __init__(self, operation):
self.operation = operation
mock_client.retrieve_credentials.return_value = DummyCall(op)
return op
@pytest.fixture
def context():
context = Mock(spec=CallbackContext)
context.user_id = "user"
context.function_call_id = "call_123"
session = Mock(spec=Session)
session.events = []
context.session = session
return context
@patch.dict(_iam_connector_credentials_provider.os.environ, clear=True)
@patch.object(_iam_connector_credentials_provider, "Client")
def test_get_client_uses_rest_transport(mock_client_class):
provider = (
_iam_connector_credentials_provider._IamConnectorCredentialsProvider()
)
provider._get_client()
mock_client_class.assert_called_once()
_, kwargs = mock_client_class.call_args
assert kwargs.get("transport") == "rest"
@patch.dict(
_iam_connector_credentials_provider.os.environ,
{"IAM_CONNECTOR_CREDENTIALS_TARGET_HOST": "some-host"},
)
@patch.object(_iam_connector_credentials_provider, "Client")
@patch.object(_iam_connector_credentials_provider, "ClientOptions")
def test_get_client_with_env_var(mock_client_options_class, mock_client_class):
provider = (
_iam_connector_credentials_provider._IamConnectorCredentialsProvider()
)
client = provider._get_client()
assert client == mock_client_class.return_value
mock_client_options_class.assert_called_once_with(api_endpoint="some-host")
mock_client_class.assert_called_once_with(
client_options=mock_client_options_class.return_value, transport="rest"
)
# ==============================================================================
# Non-interactive auth flows (API key and 2-legged OAuth)
# ==============================================================================
async def test_get_auth_credential_raises_error_if_context_is_missing(
provider, auth_scheme
):
"""Test get_auth_credential raises ValueError if context is missing."""
with pytest.raises(
ValueError,
match="GcpAuthProvider requires a context with a valid user_id",
):
await provider.get_auth_credential(auth_scheme, context=None)
async def test_get_auth_credential_raises_error_if_user_id_is_missing(
provider, auth_scheme
):
"""Test get_auth_credential raises ValueError if user_id is missing."""
context = Mock(spec=CallbackContext)
context.user_id = None
with pytest.raises(
ValueError,
match="GcpAuthProvider requires a context with a valid user_id",
):
await provider.get_auth_credential(auth_scheme, context=context)
async def test_get_auth_credential_returns_credential_if_available_immediately(
mock_client,
mock_operation,
auth_scheme,
context,
provider,
):
"""Test get_auth_credential returns credential if available immediately."""
mock_credential = RetrieveCredentialsResponse(
header="Authorization: Bearer", token="test-token"
)
mock_operation.response.value = RetrieveCredentialsResponse.serialize(
mock_credential
)
auth_credential = await provider.get_auth_credential(auth_scheme, context)
assert auth_credential.auth_type == AuthCredentialTypes.HTTP
assert auth_credential.http.scheme == "Bearer"
assert auth_credential.http.credentials.token == "test-token"
mock_client.retrieve_credentials.assert_called_once()
async def test_get_auth_credential_raises_error_if_upstream_returns_empty_header(
mock_operation,
auth_scheme,
context,
provider,
):
"""Test get_auth_credential raises RuntimeError for empty header."""
mock_credential = RetrieveCredentialsResponse(header="", token="test-token")
mock_operation.response.value = RetrieveCredentialsResponse.serialize(
mock_credential
)
with pytest.raises(
ValueError,
match=(
"Received either empty header or token from IAM Connector"
" Credentials service."
),
):
await provider.get_auth_credential(auth_scheme, context)
async def test_get_auth_credential_raises_error_if_upstream_returns_empty_token(
mock_operation,
auth_scheme,
context,
provider,
):
"""Test get_auth_credential raises RuntimeError for empty token."""
mock_credential = RetrieveCredentialsResponse(
header="Authorization: Bearer", token=""
)
mock_operation.response.value = RetrieveCredentialsResponse.serialize(
mock_credential
)
with pytest.raises(
ValueError,
match=(
"Received either empty header or token from IAM Connector"
" Credentials service."
),
):
await provider.get_auth_credential(auth_scheme, context)
async def test_get_auth_credential_returns_credential_if_upstream_returns_custom_header(
mock_operation,
auth_scheme,
context,
provider,
):
"""Test get_auth_credential returns valid credential for custom header and sets X-GOOG-API-KEY header."""
mock_credential = RetrieveCredentialsResponse(
header="some-x-api-key", token="test-token"
)
mock_operation.response.value = RetrieveCredentialsResponse.serialize(
mock_credential
)
auth_credential = await provider.get_auth_credential(auth_scheme, context)
assert auth_credential.auth_type == AuthCredentialTypes.HTTP
assert not auth_credential.http.scheme
assert auth_credential.http.credentials.token is None
assert auth_credential.http.additional_headers == {
"some-x-api-key": "test-token",
"X-GOOG-API-KEY": "test-token",
}
async def test_get_auth_credential_raises_error_if_upstream_operation_errors(
mock_operation, auth_scheme, context, provider
):
"""Test get_auth_credential raises RuntimeError for failed operations."""
mock_operation.error.message = "OAuth server error"
mock_operation.done = False
with pytest.raises(
RuntimeError, match="Operation failed: OAuth server error"
):
await provider.get_auth_credential(auth_scheme, context)
async def test_get_auth_credential_raises_error_if_upstream_call_fails(
mock_client, auth_scheme, context, provider
):
"""Test get_auth_credential raises RuntimeError for failed calls."""
mock_client.retrieve_credentials.side_effect = Exception(
"API Quota Exhausted"
)
with pytest.raises(
RuntimeError,
match="Failed to retrieve credential for user 'user' on connector",
) as exc_info:
await provider.get_auth_credential(auth_scheme, context)
# Assert that the original Exception is the chained cause!
assert str(exc_info.value.__cause__) == "API Quota Exhausted"
@patch.object(_iam_connector_credentials_provider.time, "time")
async def test_get_auth_credential_raises_error_if_polling_times_out(
mock_time,
mock_operation,
auth_scheme,
context,
provider,
):
"""Test get_auth_credential raises RuntimeError if polling times out."""
# Force the operation into the polling loop state
meta_pb = RetrieveCredentialsMetadata.pb()()
meta_pb.consent_pending.SetInParent()
meta = RetrieveCredentialsMetadata.deserialize(meta_pb.SerializeToString())
mock_operation.metadata.value = RetrieveCredentialsMetadata.serialize(meta)
# First call sets start_time=0.0, second call checks time > timeout
# (20.0 > 10.0)
mock_time.side_effect = [0.0, 20.0]
mock_metadata = Mock(spec=RetrieveCredentialsMetadata)
mock_metadata.consent_pending = True
mock_metadata.uri_consent_required = False
mock_operation.done = True
mock_operation.ClearField("error")
mock_client = Mock(spec=Client)
mock_client.retrieve_credentials.side_effect = Exception(
"Timeout waiting for credentials."
)
provider._client = mock_client
with pytest.raises(
RuntimeError,
match="Failed to retrieve credential for user 'user' on connector",
) as exc_info:
await provider.get_auth_credential(auth_scheme, context)
assert "Timeout waiting for credentials." in str(exc_info.value.__cause__)
# ==============================================================================
# Interactive Auth Flows (3-legged OAuth for User Consents)
# ==============================================================================
async def test_get_auth_credential_initiates_user_consent(
mock_operation, auth_scheme, context, provider
):
# Explicitly set the mock behavior for this test
expected_uri = "https://example.com/auth"
expected_nonce = "sample-nonce-123"
meta = RetrieveCredentialsMetadata({
"uri_consent_required": {
"authorization_uri": expected_uri,
"consent_nonce": expected_nonce,
}
})
mock_operation.metadata.value = RetrieveCredentialsMetadata.serialize(meta)
mock_operation.done = False
# Assert that there is no prior user consent completion event
assert not context.session.events
credential = await provider.get_auth_credential(auth_scheme, context)
assert credential is not None
assert credential.auth_type == AuthCredentialTypes.OAUTH2
assert credential.oauth2.auth_uri == expected_uri
assert credential.oauth2.nonce == expected_nonce
async def test_get_auth_credential_returns_fresh_auth_uri_for_repeated_requests(
mock_client, mock_operation, auth_scheme, context, provider
):
"""Test that repeated calls fetch fresh auth URIs if consent is still pending."""
# Arrange: Explicit initial URI
initial_uri = "https://example.com/auth"
initial_nonce = "initial-nonce-123"
meta1 = RetrieveCredentialsMetadata({
"uri_consent_required": {
"authorization_uri": initial_uri,
"consent_nonce": initial_nonce,
}
})
mock_operation.metadata.value = RetrieveCredentialsMetadata.serialize(meta1)
mock_operation.done = False
credential1 = await provider.get_auth_credential(auth_scheme, context)
assert credential1.oauth2.auth_uri == initial_uri
assert credential1.oauth2.nonce == initial_nonce
# Arrange: Explicit new URI for the second call
fresh_auth_uri = "https://example.com/auth_new"
fresh_nonce = "fresh-nonce-456"
meta2 = RetrieveCredentialsMetadata({
"uri_consent_required": {
"authorization_uri": fresh_auth_uri,
"consent_nonce": fresh_nonce,
}
})
mock_operation.metadata.value = RetrieveCredentialsMetadata.serialize(meta2)
credential2 = await provider.get_auth_credential(auth_scheme, context)
assert mock_client.retrieve_credentials.call_count == 2
assert credential2.oauth2.auth_uri == fresh_auth_uri
assert credential2.oauth2.nonce == fresh_nonce
async def test_get_auth_credential_returns_token_if_consent_was_completed(
mock_operation, auth_scheme, context, provider
):
# Setup mock credential for successful credential retrieval
mock_credential = RetrieveCredentialsResponse(
header="Authorization: Bearer", token="test-token"
)
mock_operation.response.value = RetrieveCredentialsResponse.serialize(
mock_credential
)
# Create mock events
# 1. FunctionCall event for adk_request_credential
function_call = Mock()
function_call.id = "auth-req-1"
function_call.name = REQUEST_EUC_FUNCTION_CALL_NAME
function_call.args = AuthToolArguments(
function_call_id="call-123",
auth_config=Mock(spec=AuthConfig, auth_scheme=auth_scheme),
).model_dump(by_alias=True, exclude_none=True)
event1 = Mock()
event1.get_function_calls.return_value = [function_call]
event1.get_function_responses.return_value = []
# 2. FunctionResponse event for adk_request_credential
function_response = Mock()
function_response.id = "auth-req-1"
function_response.name = REQUEST_EUC_FUNCTION_CALL_NAME
event2 = Mock()
event2.get_function_calls.return_value = []
event2.get_function_responses.return_value = [function_response]
# Setup tool context and event history (order of events matters)
context.session.events = [event1, event2]
context.function_call_id = "call-123"
# Also set uri_consent_required to True-ish so it enters the check block
meta = RetrieveCredentialsMetadata(
uri_consent_required=RetrieveCredentialsMetadata.UriConsentRequired()
)
mock_operation.metadata.value = RetrieveCredentialsMetadata.serialize(meta)
# Execute
auth_credential = await provider.get_auth_credential(auth_scheme, context)
# Verify
assert auth_credential is not None
assert auth_credential.auth_type == AuthCredentialTypes.HTTP
assert auth_credential.http.scheme == "Bearer"
assert auth_credential.http.credentials.token == "test-token"
async def test_get_auth_credential_raises_error_if_consent_canceled(
mock_operation, auth_scheme, context, provider
):
function_call = Mock()
function_call.id = "auth-req-1"
function_call.name = REQUEST_EUC_FUNCTION_CALL_NAME
function_call.args = AuthToolArguments(
function_call_id="call-123",
auth_config=Mock(spec=AuthConfig, auth_scheme=auth_scheme),
).model_dump(by_alias=True, exclude_none=True)
event1 = Mock()
event1.get_function_calls.return_value = [function_call]
event1.get_function_responses.return_value = []
function_response = Mock()
function_response.id = "auth-req-1"
function_response.name = REQUEST_EUC_FUNCTION_CALL_NAME
event2 = Mock()
event2.get_function_calls.return_value = []
event2.get_function_responses.return_value = [function_response]
context.session.events = [event1, event2]
context.function_call_id = "call-123"
meta = RetrieveCredentialsMetadata({
"uri_consent_required": {
"authorization_uri": "https://example.com/auth",
"consent_nonce": "sample-nonce",
}
})
mock_operation.metadata.value = RetrieveCredentialsMetadata.serialize(meta)
mock_operation.done = False
with pytest.raises(
RuntimeError, match="Failed to retrieve consent based credential."
):
await provider.get_auth_credential(auth_scheme, context)
async def test_get_auth_credential_handles_consent_pending_state_correctly(
mock_client,
auth_scheme,
context,
provider,
):
"""Test get_auth_credential enters the polling loop when consent is pending."""
# 1. Setup the first retrieve_credentials call to return a pending Operation
# containing consent_pending metadata.
meta_pb = RetrieveCredentialsMetadata.pb()()
meta_pb.consent_pending.SetInParent()
op_pending = Operation(done=False)
op_pending.metadata.value = meta_pb.SerializeToString()
# 2. Setup the second retrieve_credentials call (the poll) to return success.
op_success = Operation(done=True)
resp = RetrieveCredentialsResponse(
header="Authorization: Bearer", token="valid-token"
)
op_success.response.value = RetrieveCredentialsResponse.serialize(resp)
# Configure mock client to return pending first, then success
mock_client.retrieve_credentials.side_effect = [
Mock(operation=op_pending),
Mock(operation=op_success),
]
# 3. Call the provider
credential = await provider.get_auth_credential(auth_scheme, context)
# 4. Verify that it polled and successfully returned the token
assert credential is not None
assert credential.http.credentials.token == "valid-token"
# Verify that retrieve_credentials was called twice (initial + 1 poll)
assert mock_client.retrieve_credentials.call_count == 2