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
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:
+423
@@ -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
|
||||
)
|
||||
+510
@@ -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
|
||||
Reference in New Issue
Block a user