Skip to content

Commit 4cc210e

Browse files
authored
Merge branch 'main' into fix
2 parents 71e72ff + a614507 commit 4cc210e

2 files changed

Lines changed: 40 additions & 1 deletion

File tree

src/google/adk/integrations/agent_identity/_iam_connector_credentials_provider.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -238,7 +238,7 @@ async def get_auth_credential(
238238
logger.debug("Auth credential obtained immediately.")
239239
return _construct_auth_credential(response)
240240

241-
if metadata and metadata.consent_pending:
241+
if metadata is not None and "consent_pending" in metadata:
242242
# Get 2-legged OAuth token. Allow enough time for token exchange.
243243
try:
244244
operation = await self._poll_credentials(

tests/unittests/integrations/agent_identity/test_iam_connector_credentials_provider.py

Lines changed: 39 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -469,3 +469,42 @@ async def test_get_auth_credential_raises_error_if_consent_canceled(
469469
RuntimeError, match="Failed to retrieve consent based credential."
470470
):
471471
await provider.get_auth_credential(auth_scheme, context)
472+
473+
474+
async def test_get_auth_credential_handles_consent_pending_state_correctly(
475+
mock_client,
476+
auth_scheme,
477+
context,
478+
provider,
479+
):
480+
"""Test get_auth_credential enters the polling loop when consent is pending."""
481+
# 1. Setup the first retrieve_credentials call to return a pending Operation
482+
# containing consent_pending metadata.
483+
meta_pb = RetrieveCredentialsMetadata.pb()()
484+
meta_pb.consent_pending.SetInParent()
485+
486+
op_pending = Operation(done=False)
487+
op_pending.metadata.value = meta_pb.SerializeToString()
488+
489+
# 2. Setup the second retrieve_credentials call (the poll) to return success.
490+
op_success = Operation(done=True)
491+
resp = RetrieveCredentialsResponse(
492+
header="Authorization: Bearer", token="valid-token"
493+
)
494+
op_success.response.value = RetrieveCredentialsResponse.serialize(resp)
495+
496+
# Configure mock client to return pending first, then success
497+
mock_client.retrieve_credentials.side_effect = [
498+
Mock(operation=op_pending),
499+
Mock(operation=op_success),
500+
]
501+
502+
# 3. Call the provider
503+
credential = await provider.get_auth_credential(auth_scheme, context)
504+
505+
# 4. Verify that it polled and successfully returned the token
506+
assert credential is not None
507+
assert credential.http.credentials.token == "valid-token"
508+
509+
# Verify that retrieve_credentials was called twice (initial + 1 poll)
510+
assert mock_client.retrieve_credentials.call_count == 2

0 commit comments

Comments
 (0)