mirror of
https://github.com/encounter/adk-python.git
synced 2026-07-09 18:19:28 -07:00
refactor: Refactor oauth2_credential_exchanger to exchanger and refresher separately
PiperOrigin-RevId: 772979993
This commit is contained in:
Xiang (Sean) Zhou
committed by
Copybara-Service
Copybara-Service
parent
a17ebe6ebd
commit
9a207cb832
@@ -0,0 +1,220 @@
|
||||
# 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.
|
||||
|
||||
import time
|
||||
from unittest.mock import Mock
|
||||
from unittest.mock import patch
|
||||
|
||||
from authlib.oauth2.rfc6749 import OAuth2Token
|
||||
from google.adk.auth.auth_credential import AuthCredential
|
||||
from google.adk.auth.auth_credential import AuthCredentialTypes
|
||||
from google.adk.auth.auth_credential import OAuth2Auth
|
||||
from google.adk.auth.auth_schemes import OpenIdConnectWithConfig
|
||||
from google.adk.auth.exchanger.base_credential_exchanger import CredentialExchangError
|
||||
from google.adk.auth.exchanger.oauth2_credential_exchanger import OAuth2CredentialExchanger
|
||||
import pytest
|
||||
|
||||
|
||||
class TestOAuth2CredentialExchanger:
|
||||
"""Test suite for OAuth2CredentialExchanger."""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_exchange_with_existing_token(self):
|
||||
"""Test exchange method when access token already exists."""
|
||||
scheme = OpenIdConnectWithConfig(
|
||||
type_="openIdConnect",
|
||||
openId_connect_url=(
|
||||
"https://example.com/.well-known/openid_configuration"
|
||||
),
|
||||
authorization_endpoint="https://example.com/auth",
|
||||
token_endpoint="https://example.com/token",
|
||||
scopes=["openid"],
|
||||
)
|
||||
credential = AuthCredential(
|
||||
auth_type=AuthCredentialTypes.OPEN_ID_CONNECT,
|
||||
oauth2=OAuth2Auth(
|
||||
client_id="test_client_id",
|
||||
client_secret="test_client_secret",
|
||||
access_token="existing_token",
|
||||
),
|
||||
)
|
||||
|
||||
exchanger = OAuth2CredentialExchanger()
|
||||
result = await exchanger.exchange(credential, scheme)
|
||||
|
||||
# Should return the same credential since access token already exists
|
||||
assert result == credential
|
||||
assert result.oauth2.access_token == "existing_token"
|
||||
|
||||
@patch("google.adk.auth.oauth2_credential_util.OAuth2Session")
|
||||
@pytest.mark.asyncio
|
||||
async def test_exchange_success(self, mock_oauth2_session):
|
||||
"""Test successful token exchange."""
|
||||
# Setup mock
|
||||
mock_client = Mock()
|
||||
mock_oauth2_session.return_value = mock_client
|
||||
mock_tokens = OAuth2Token({
|
||||
"access_token": "new_access_token",
|
||||
"refresh_token": "new_refresh_token",
|
||||
"expires_at": int(time.time()) + 3600,
|
||||
"expires_in": 3600,
|
||||
})
|
||||
mock_client.fetch_token.return_value = mock_tokens
|
||||
|
||||
scheme = OpenIdConnectWithConfig(
|
||||
type_="openIdConnect",
|
||||
openId_connect_url=(
|
||||
"https://example.com/.well-known/openid_configuration"
|
||||
),
|
||||
authorization_endpoint="https://example.com/auth",
|
||||
token_endpoint="https://example.com/token",
|
||||
scopes=["openid"],
|
||||
)
|
||||
credential = AuthCredential(
|
||||
auth_type=AuthCredentialTypes.OPEN_ID_CONNECT,
|
||||
oauth2=OAuth2Auth(
|
||||
client_id="test_client_id",
|
||||
client_secret="test_client_secret",
|
||||
auth_response_uri="https://example.com/callback?code=auth_code",
|
||||
auth_code="auth_code",
|
||||
),
|
||||
)
|
||||
|
||||
exchanger = OAuth2CredentialExchanger()
|
||||
result = await exchanger.exchange(credential, scheme)
|
||||
|
||||
# Verify token exchange was successful
|
||||
assert result.oauth2.access_token == "new_access_token"
|
||||
assert result.oauth2.refresh_token == "new_refresh_token"
|
||||
mock_client.fetch_token.assert_called_once()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_exchange_missing_auth_scheme(self):
|
||||
"""Test exchange with missing auth_scheme raises ValueError."""
|
||||
credential = AuthCredential(
|
||||
auth_type=AuthCredentialTypes.OPEN_ID_CONNECT,
|
||||
oauth2=OAuth2Auth(
|
||||
client_id="test_client_id",
|
||||
client_secret="test_client_secret",
|
||||
),
|
||||
)
|
||||
|
||||
exchanger = OAuth2CredentialExchanger()
|
||||
try:
|
||||
await exchanger.exchange(credential, None)
|
||||
assert False, "Should have raised ValueError"
|
||||
except CredentialExchangError as e:
|
||||
assert "auth_scheme is required" in str(e)
|
||||
|
||||
@patch("google.adk.auth.oauth2_credential_util.OAuth2Session")
|
||||
@pytest.mark.asyncio
|
||||
async def test_exchange_no_session(self, mock_oauth2_session):
|
||||
"""Test exchange when OAuth2Session cannot be created."""
|
||||
# Mock to return None for create_oauth2_session
|
||||
mock_oauth2_session.return_value = None
|
||||
|
||||
scheme = OpenIdConnectWithConfig(
|
||||
type_="openIdConnect",
|
||||
openId_connect_url=(
|
||||
"https://example.com/.well-known/openid_configuration"
|
||||
),
|
||||
authorization_endpoint="https://example.com/auth",
|
||||
token_endpoint="https://example.com/token",
|
||||
scopes=["openid"],
|
||||
)
|
||||
credential = AuthCredential(
|
||||
auth_type=AuthCredentialTypes.OPEN_ID_CONNECT,
|
||||
oauth2=OAuth2Auth(
|
||||
client_id="test_client_id",
|
||||
# Missing client_secret to trigger session creation failure
|
||||
),
|
||||
)
|
||||
|
||||
exchanger = OAuth2CredentialExchanger()
|
||||
result = await exchanger.exchange(credential, scheme)
|
||||
|
||||
# Should return original credential when session creation fails
|
||||
assert result == credential
|
||||
assert result.oauth2.access_token is None
|
||||
|
||||
@patch("google.adk.auth.oauth2_credential_util.OAuth2Session")
|
||||
@pytest.mark.asyncio
|
||||
async def test_exchange_fetch_token_failure(self, mock_oauth2_session):
|
||||
"""Test exchange when fetch_token fails."""
|
||||
# Setup mock to raise exception during fetch_token
|
||||
mock_client = Mock()
|
||||
mock_oauth2_session.return_value = mock_client
|
||||
mock_client.fetch_token.side_effect = Exception("Token fetch failed")
|
||||
|
||||
scheme = OpenIdConnectWithConfig(
|
||||
type_="openIdConnect",
|
||||
openId_connect_url=(
|
||||
"https://example.com/.well-known/openid_configuration"
|
||||
),
|
||||
authorization_endpoint="https://example.com/auth",
|
||||
token_endpoint="https://example.com/token",
|
||||
scopes=["openid"],
|
||||
)
|
||||
credential = AuthCredential(
|
||||
auth_type=AuthCredentialTypes.OPEN_ID_CONNECT,
|
||||
oauth2=OAuth2Auth(
|
||||
client_id="test_client_id",
|
||||
client_secret="test_client_secret",
|
||||
auth_response_uri="https://example.com/callback?code=auth_code",
|
||||
auth_code="auth_code",
|
||||
),
|
||||
)
|
||||
|
||||
exchanger = OAuth2CredentialExchanger()
|
||||
result = await exchanger.exchange(credential, scheme)
|
||||
|
||||
# Should return original credential when fetch_token fails
|
||||
assert result == credential
|
||||
assert result.oauth2.access_token is None
|
||||
mock_client.fetch_token.assert_called_once()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_exchange_authlib_not_available(self):
|
||||
"""Test exchange when authlib is not available."""
|
||||
scheme = OpenIdConnectWithConfig(
|
||||
type_="openIdConnect",
|
||||
openId_connect_url=(
|
||||
"https://example.com/.well-known/openid_configuration"
|
||||
),
|
||||
authorization_endpoint="https://example.com/auth",
|
||||
token_endpoint="https://example.com/token",
|
||||
scopes=["openid"],
|
||||
)
|
||||
credential = AuthCredential(
|
||||
auth_type=AuthCredentialTypes.OPEN_ID_CONNECT,
|
||||
oauth2=OAuth2Auth(
|
||||
client_id="test_client_id",
|
||||
client_secret="test_client_secret",
|
||||
auth_response_uri="https://example.com/callback?code=auth_code",
|
||||
auth_code="auth_code",
|
||||
),
|
||||
)
|
||||
|
||||
exchanger = OAuth2CredentialExchanger()
|
||||
|
||||
# Mock AUTHLIB_AVIALABLE to False
|
||||
with patch(
|
||||
"google.adk.auth.exchanger.oauth2_credential_exchanger.AUTHLIB_AVIALABLE",
|
||||
False,
|
||||
):
|
||||
result = await exchanger.exchange(credential, scheme)
|
||||
|
||||
# Should return original credential when authlib is not available
|
||||
assert result == credential
|
||||
assert result.oauth2.access_token is None
|
||||
@@ -0,0 +1,13 @@
|
||||
# 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.
|
||||
@@ -0,0 +1,297 @@
|
||||
# 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.
|
||||
|
||||
import time
|
||||
from unittest.mock import Mock
|
||||
from unittest.mock import patch
|
||||
|
||||
from authlib.oauth2.rfc6749 import OAuth2Token
|
||||
from google.adk.auth.auth_credential import AuthCredential
|
||||
from google.adk.auth.auth_credential import AuthCredentialTypes
|
||||
from google.adk.auth.auth_credential import OAuth2Auth
|
||||
from google.adk.auth.auth_schemes import OpenIdConnectWithConfig
|
||||
from google.adk.auth.refresher.oauth2_credential_refresher import OAuth2CredentialRefresher
|
||||
import pytest
|
||||
|
||||
|
||||
class TestOAuth2CredentialRefresher:
|
||||
"""Test suite for OAuth2CredentialRefresher."""
|
||||
|
||||
@patch("google.adk.auth.refresher.oauth2_credential_refresher.OAuth2Token")
|
||||
@pytest.mark.asyncio
|
||||
async def test_needs_refresh_token_not_expired(self, mock_oauth2_token):
|
||||
"""Test needs_refresh when token is not expired."""
|
||||
mock_token_instance = Mock()
|
||||
mock_token_instance.is_expired.return_value = False
|
||||
mock_oauth2_token.return_value = mock_token_instance
|
||||
|
||||
scheme = OpenIdConnectWithConfig(
|
||||
type_="openIdConnect",
|
||||
openId_connect_url=(
|
||||
"https://example.com/.well-known/openid_configuration"
|
||||
),
|
||||
authorization_endpoint="https://example.com/auth",
|
||||
token_endpoint="https://example.com/token",
|
||||
scopes=["openid"],
|
||||
)
|
||||
credential = AuthCredential(
|
||||
auth_type=AuthCredentialTypes.OPEN_ID_CONNECT,
|
||||
oauth2=OAuth2Auth(
|
||||
client_id="test_client_id",
|
||||
client_secret="test_client_secret",
|
||||
access_token="existing_token",
|
||||
expires_at=int(time.time()) + 3600,
|
||||
),
|
||||
)
|
||||
|
||||
refresher = OAuth2CredentialRefresher()
|
||||
needs_refresh = await refresher.is_refresh_needed(credential, scheme)
|
||||
|
||||
assert not needs_refresh
|
||||
|
||||
@patch("google.adk.auth.refresher.oauth2_credential_refresher.OAuth2Token")
|
||||
@pytest.mark.asyncio
|
||||
async def test_needs_refresh_token_expired(self, mock_oauth2_token):
|
||||
"""Test needs_refresh when token is expired."""
|
||||
mock_token_instance = Mock()
|
||||
mock_token_instance.is_expired.return_value = True
|
||||
mock_oauth2_token.return_value = mock_token_instance
|
||||
|
||||
scheme = OpenIdConnectWithConfig(
|
||||
type_="openIdConnect",
|
||||
openId_connect_url=(
|
||||
"https://example.com/.well-known/openid_configuration"
|
||||
),
|
||||
authorization_endpoint="https://example.com/auth",
|
||||
token_endpoint="https://example.com/token",
|
||||
scopes=["openid"],
|
||||
)
|
||||
credential = AuthCredential(
|
||||
auth_type=AuthCredentialTypes.OPEN_ID_CONNECT,
|
||||
oauth2=OAuth2Auth(
|
||||
client_id="test_client_id",
|
||||
client_secret="test_client_secret",
|
||||
access_token="existing_token",
|
||||
expires_at=int(time.time()) - 3600, # Expired
|
||||
),
|
||||
)
|
||||
|
||||
refresher = OAuth2CredentialRefresher()
|
||||
needs_refresh = await refresher.is_refresh_needed(credential, scheme)
|
||||
|
||||
assert needs_refresh
|
||||
|
||||
@patch("google.adk.auth.oauth2_credential_util.OAuth2Session")
|
||||
@patch("google.adk.auth.oauth2_credential_util.OAuth2Token")
|
||||
@pytest.mark.asyncio
|
||||
async def test_refresh_token_expired_success(
|
||||
self, mock_oauth2_token, mock_oauth2_session
|
||||
):
|
||||
"""Test successful token refresh when token is expired."""
|
||||
# Setup mock token
|
||||
mock_token_instance = Mock()
|
||||
mock_token_instance.is_expired.return_value = True
|
||||
mock_oauth2_token.return_value = mock_token_instance
|
||||
|
||||
# Setup mock session
|
||||
mock_client = Mock()
|
||||
mock_oauth2_session.return_value = mock_client
|
||||
mock_tokens = OAuth2Token({
|
||||
"access_token": "refreshed_access_token",
|
||||
"refresh_token": "refreshed_refresh_token",
|
||||
"expires_at": int(time.time()) + 3600,
|
||||
"expires_in": 3600,
|
||||
})
|
||||
mock_client.refresh_token.return_value = mock_tokens
|
||||
|
||||
scheme = OpenIdConnectWithConfig(
|
||||
type_="openIdConnect",
|
||||
openId_connect_url=(
|
||||
"https://example.com/.well-known/openid_configuration"
|
||||
),
|
||||
authorization_endpoint="https://example.com/auth",
|
||||
token_endpoint="https://example.com/token",
|
||||
scopes=["openid"],
|
||||
)
|
||||
credential = AuthCredential(
|
||||
auth_type=AuthCredentialTypes.OPEN_ID_CONNECT,
|
||||
oauth2=OAuth2Auth(
|
||||
client_id="test_client_id",
|
||||
client_secret="test_client_secret",
|
||||
access_token="old_token",
|
||||
refresh_token="old_refresh_token",
|
||||
expires_at=int(time.time()) - 3600, # Expired
|
||||
),
|
||||
)
|
||||
|
||||
refresher = OAuth2CredentialRefresher()
|
||||
result = await refresher.refresh(credential, scheme)
|
||||
|
||||
# Verify token refresh was successful
|
||||
assert result.oauth2.access_token == "refreshed_access_token"
|
||||
assert result.oauth2.refresh_token == "refreshed_refresh_token"
|
||||
mock_client.refresh_token.assert_called_once()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_refresh_no_oauth2_credential(self):
|
||||
"""Test refresh with no OAuth2 credential returns original."""
|
||||
scheme = OpenIdConnectWithConfig(
|
||||
type_="openIdConnect",
|
||||
openId_connect_url=(
|
||||
"https://example.com/.well-known/openid_configuration"
|
||||
),
|
||||
authorization_endpoint="https://example.com/auth",
|
||||
token_endpoint="https://example.com/token",
|
||||
scopes=["openid"],
|
||||
)
|
||||
credential = AuthCredential(
|
||||
auth_type=AuthCredentialTypes.OPEN_ID_CONNECT,
|
||||
# No oauth2 field
|
||||
)
|
||||
|
||||
refresher = OAuth2CredentialRefresher()
|
||||
result = await refresher.refresh(credential, scheme)
|
||||
|
||||
assert result == credential
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_needs_refresh_google_oauth2_json_expired(self):
|
||||
"""Test needs_refresh with Google OAuth2 JSON credential that is expired."""
|
||||
import json
|
||||
from unittest.mock import patch
|
||||
|
||||
# Mock Google OAuth2 JSON credential data
|
||||
google_oauth2_json = json.dumps({
|
||||
"client_id": "test_client_id",
|
||||
"client_secret": "test_client_secret",
|
||||
"refresh_token": "test_refresh_token",
|
||||
"type": "authorized_user",
|
||||
})
|
||||
|
||||
credential = AuthCredential(
|
||||
auth_type=AuthCredentialTypes.OAUTH2,
|
||||
google_oauth2_json=google_oauth2_json,
|
||||
)
|
||||
|
||||
# Mock the Google Credentials class
|
||||
with patch(
|
||||
"google.adk.auth.refresher.oauth2_credential_refresher.Credentials"
|
||||
) as mock_credentials:
|
||||
mock_google_credential = Mock()
|
||||
mock_google_credential.expired = True
|
||||
mock_google_credential.refresh_token = "test_refresh_token"
|
||||
mock_credentials.from_authorized_user_info.return_value = (
|
||||
mock_google_credential
|
||||
)
|
||||
|
||||
refresher = OAuth2CredentialRefresher()
|
||||
needs_refresh = await refresher.is_refresh_needed(credential, None)
|
||||
|
||||
assert needs_refresh
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_needs_refresh_google_oauth2_json_not_expired(self):
|
||||
"""Test needs_refresh with Google OAuth2 JSON credential that is not expired."""
|
||||
import json
|
||||
from unittest.mock import patch
|
||||
|
||||
# Mock Google OAuth2 JSON credential data
|
||||
google_oauth2_json = json.dumps({
|
||||
"client_id": "test_client_id",
|
||||
"client_secret": "test_client_secret",
|
||||
"refresh_token": "test_refresh_token",
|
||||
"type": "authorized_user",
|
||||
})
|
||||
|
||||
credential = AuthCredential(
|
||||
auth_type=AuthCredentialTypes.OAUTH2,
|
||||
google_oauth2_json=google_oauth2_json,
|
||||
)
|
||||
|
||||
# Mock the Google Credentials class
|
||||
with patch(
|
||||
"google.adk.auth.refresher.oauth2_credential_refresher.Credentials"
|
||||
) as mock_credentials:
|
||||
mock_google_credential = Mock()
|
||||
mock_google_credential.expired = False
|
||||
mock_google_credential.refresh_token = "test_refresh_token"
|
||||
mock_credentials.from_authorized_user_info.return_value = (
|
||||
mock_google_credential
|
||||
)
|
||||
|
||||
refresher = OAuth2CredentialRefresher()
|
||||
needs_refresh = await refresher.is_refresh_needed(credential, None)
|
||||
|
||||
assert not needs_refresh
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_refresh_google_oauth2_json_success(self):
|
||||
"""Test successful refresh of Google OAuth2 JSON credential."""
|
||||
import json
|
||||
from unittest.mock import patch
|
||||
|
||||
# Mock Google OAuth2 JSON credential data
|
||||
google_oauth2_json = json.dumps({
|
||||
"client_id": "test_client_id",
|
||||
"client_secret": "test_client_secret",
|
||||
"refresh_token": "test_refresh_token",
|
||||
"type": "authorized_user",
|
||||
})
|
||||
|
||||
credential = AuthCredential(
|
||||
auth_type=AuthCredentialTypes.OAUTH2,
|
||||
google_oauth2_json=google_oauth2_json,
|
||||
)
|
||||
|
||||
# Mock the Google Credentials and Request classes
|
||||
with patch(
|
||||
"google.adk.auth.refresher.oauth2_credential_refresher.Credentials"
|
||||
) as mock_credentials:
|
||||
with patch(
|
||||
"google.adk.auth.refresher.oauth2_credential_refresher.Request"
|
||||
) as mock_request:
|
||||
mock_google_credential = Mock()
|
||||
mock_google_credential.expired = True
|
||||
mock_google_credential.refresh_token = "test_refresh_token"
|
||||
mock_google_credential.to_json.return_value = json.dumps({
|
||||
"client_id": "test_client_id",
|
||||
"client_secret": "test_client_secret",
|
||||
"refresh_token": "new_refresh_token",
|
||||
"access_token": "new_access_token",
|
||||
"type": "authorized_user",
|
||||
})
|
||||
mock_credentials.from_authorized_user_info.return_value = (
|
||||
mock_google_credential
|
||||
)
|
||||
|
||||
refresher = OAuth2CredentialRefresher()
|
||||
result = await refresher.refresh(credential, None)
|
||||
|
||||
mock_google_credential.refresh.assert_called_once()
|
||||
assert (
|
||||
result.google_oauth2_json != google_oauth2_json
|
||||
) # Should be updated
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_needs_refresh_no_oauth2_credential(self):
|
||||
"""Test needs_refresh with no OAuth2 credential returns False."""
|
||||
credential = AuthCredential(
|
||||
auth_type=AuthCredentialTypes.HTTP,
|
||||
# No oauth2 field
|
||||
)
|
||||
|
||||
refresher = OAuth2CredentialRefresher()
|
||||
needs_refresh = await refresher.is_refresh_needed(credential, None)
|
||||
|
||||
assert not needs_refresh
|
||||
@@ -13,8 +13,11 @@
|
||||
# limitations under the License.
|
||||
|
||||
import copy
|
||||
import time
|
||||
from unittest.mock import Mock
|
||||
from unittest.mock import patch
|
||||
|
||||
from authlib.oauth2.rfc6749 import OAuth2Token
|
||||
from fastapi.openapi.models import APIKey
|
||||
from fastapi.openapi.models import APIKeyIn
|
||||
from fastapi.openapi.models import OAuth2
|
||||
@@ -405,7 +408,8 @@ class TestGetAuthResponse:
|
||||
class TestParseAndStoreAuthResponse:
|
||||
"""Tests for the parse_and_store_auth_response method."""
|
||||
|
||||
def test_non_oauth_scheme(self, auth_config_with_exchanged):
|
||||
@pytest.mark.asyncio
|
||||
async def test_non_oauth_scheme(self, auth_config_with_exchanged):
|
||||
"""Test with a non-OAuth auth scheme."""
|
||||
# Modify the auth scheme type to be non-OAuth
|
||||
auth_config = copy.deepcopy(auth_config_with_exchanged)
|
||||
@@ -416,7 +420,7 @@ class TestParseAndStoreAuthResponse:
|
||||
handler = AuthHandler(auth_config)
|
||||
state = MockState()
|
||||
|
||||
handler.parse_and_store_auth_response(state)
|
||||
await handler.parse_and_store_auth_response(state)
|
||||
|
||||
credential_key = auth_config.credential_key
|
||||
assert (
|
||||
@@ -424,7 +428,10 @@ class TestParseAndStoreAuthResponse:
|
||||
)
|
||||
|
||||
@patch("google.adk.auth.auth_handler.AuthHandler.exchange_auth_token")
|
||||
def test_oauth_scheme(self, mock_exchange_token, auth_config_with_exchanged):
|
||||
@pytest.mark.asyncio
|
||||
async def test_oauth_scheme(
|
||||
self, mock_exchange_token, auth_config_with_exchanged
|
||||
):
|
||||
"""Test with an OAuth auth scheme."""
|
||||
mock_exchange_token.return_value = AuthCredential(
|
||||
auth_type=AuthCredentialTypes.OAUTH2,
|
||||
@@ -434,7 +441,7 @@ class TestParseAndStoreAuthResponse:
|
||||
handler = AuthHandler(auth_config_with_exchanged)
|
||||
state = MockState()
|
||||
|
||||
handler.parse_and_store_auth_response(state)
|
||||
await handler.parse_and_store_auth_response(state)
|
||||
|
||||
credential_key = auth_config_with_exchanged.credential_key
|
||||
assert state["temp:" + credential_key] == mock_exchange_token.return_value
|
||||
@@ -444,20 +451,20 @@ class TestParseAndStoreAuthResponse:
|
||||
class TestExchangeAuthToken:
|
||||
"""Tests for the exchange_auth_token method."""
|
||||
|
||||
def test_token_exchange_not_supported(
|
||||
@pytest.mark.asyncio
|
||||
async def test_token_exchange_not_supported(
|
||||
self, auth_config_with_auth_code, monkeypatch
|
||||
):
|
||||
"""Test when token exchange is not supported."""
|
||||
monkeypatch.setattr(
|
||||
"google.adk.auth.oauth2_credential_fetcher.AUTHLIB_AVIALABLE", False
|
||||
)
|
||||
monkeypatch.setattr("google.adk.auth.auth_handler.AUTHLIB_AVIALABLE", False)
|
||||
|
||||
handler = AuthHandler(auth_config_with_auth_code)
|
||||
result = handler.exchange_auth_token()
|
||||
result = await handler.exchange_auth_token()
|
||||
|
||||
assert result == auth_config_with_auth_code.exchanged_auth_credential
|
||||
|
||||
def test_openid_missing_token_endpoint(
|
||||
@pytest.mark.asyncio
|
||||
async def test_openid_missing_token_endpoint(
|
||||
self, openid_auth_scheme, oauth2_credentials_with_auth_code
|
||||
):
|
||||
"""Test OpenID Connect without a token endpoint."""
|
||||
@@ -472,11 +479,12 @@ class TestExchangeAuthToken:
|
||||
)
|
||||
|
||||
handler = AuthHandler(config)
|
||||
result = handler.exchange_auth_token()
|
||||
result = await handler.exchange_auth_token()
|
||||
|
||||
assert result == oauth2_credentials_with_auth_code
|
||||
|
||||
def test_oauth2_missing_token_url(
|
||||
@pytest.mark.asyncio
|
||||
async def test_oauth2_missing_token_url(
|
||||
self, oauth2_auth_scheme, oauth2_credentials_with_auth_code
|
||||
):
|
||||
"""Test OAuth2 without a token URL."""
|
||||
@@ -491,11 +499,12 @@ class TestExchangeAuthToken:
|
||||
)
|
||||
|
||||
handler = AuthHandler(config)
|
||||
result = handler.exchange_auth_token()
|
||||
result = await handler.exchange_auth_token()
|
||||
|
||||
assert result == oauth2_credentials_with_auth_code
|
||||
|
||||
def test_non_oauth_scheme(self, auth_config_with_auth_code):
|
||||
@pytest.mark.asyncio
|
||||
async def test_non_oauth_scheme(self, auth_config_with_auth_code):
|
||||
"""Test with a non-OAuth auth scheme."""
|
||||
# Modify the auth scheme type to be non-OAuth
|
||||
auth_config = copy.deepcopy(auth_config_with_auth_code)
|
||||
@@ -504,11 +513,12 @@ class TestExchangeAuthToken:
|
||||
)
|
||||
|
||||
handler = AuthHandler(auth_config)
|
||||
result = handler.exchange_auth_token()
|
||||
result = await handler.exchange_auth_token()
|
||||
|
||||
assert result == auth_config.exchanged_auth_credential
|
||||
|
||||
def test_missing_credentials(self, oauth2_auth_scheme):
|
||||
@pytest.mark.asyncio
|
||||
async def test_missing_credentials(self, oauth2_auth_scheme):
|
||||
"""Test with missing credentials."""
|
||||
empty_credential = AuthCredential(auth_type=AuthCredentialTypes.OAUTH2)
|
||||
|
||||
@@ -518,11 +528,12 @@ class TestExchangeAuthToken:
|
||||
)
|
||||
|
||||
handler = AuthHandler(config)
|
||||
result = handler.exchange_auth_token()
|
||||
result = await handler.exchange_auth_token()
|
||||
|
||||
assert result == empty_credential
|
||||
|
||||
def test_credentials_with_token(
|
||||
@pytest.mark.asyncio
|
||||
async def test_credentials_with_token(
|
||||
self, auth_config, oauth2_credentials_with_token
|
||||
):
|
||||
"""Test when credentials already have a token."""
|
||||
@@ -533,18 +544,29 @@ class TestExchangeAuthToken:
|
||||
)
|
||||
|
||||
handler = AuthHandler(config)
|
||||
result = handler.exchange_auth_token()
|
||||
result = await handler.exchange_auth_token()
|
||||
|
||||
assert result == oauth2_credentials_with_token
|
||||
|
||||
@patch(
|
||||
"google.adk.auth.oauth2_credential_util.OAuth2Session",
|
||||
MockOAuth2Session,
|
||||
)
|
||||
def test_successful_token_exchange(self, auth_config_with_auth_code):
|
||||
@patch("google.adk.auth.oauth2_credential_util.OAuth2Session")
|
||||
@pytest.mark.asyncio
|
||||
async def test_successful_token_exchange(
|
||||
self, mock_oauth2_session, auth_config_with_auth_code
|
||||
):
|
||||
"""Test a successful token exchange."""
|
||||
# Setup mock OAuth2Session
|
||||
mock_client = Mock()
|
||||
mock_oauth2_session.return_value = mock_client
|
||||
mock_tokens = OAuth2Token({
|
||||
"access_token": "mock_access_token",
|
||||
"refresh_token": "mock_refresh_token",
|
||||
"expires_at": int(time.time()) + 3600,
|
||||
"expires_in": 3600,
|
||||
})
|
||||
mock_client.fetch_token.return_value = mock_tokens
|
||||
|
||||
handler = AuthHandler(auth_config_with_auth_code)
|
||||
result = handler.exchange_auth_token()
|
||||
result = await handler.exchange_auth_token()
|
||||
|
||||
assert result.oauth2.access_token == "mock_access_token"
|
||||
assert result.oauth2.refresh_token == "mock_refresh_token"
|
||||
|
||||
@@ -1,147 +0,0 @@
|
||||
# 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.
|
||||
|
||||
import time
|
||||
from unittest.mock import Mock
|
||||
|
||||
from authlib.oauth2.rfc6749 import OAuth2Token
|
||||
from fastapi.openapi.models import OAuth2
|
||||
from fastapi.openapi.models import OAuthFlowAuthorizationCode
|
||||
from fastapi.openapi.models import OAuthFlows
|
||||
from google.adk.auth.auth_credential import AuthCredential
|
||||
from google.adk.auth.auth_credential import AuthCredentialTypes
|
||||
from google.adk.auth.auth_credential import OAuth2Auth
|
||||
from google.adk.auth.auth_schemes import OpenIdConnectWithConfig
|
||||
from google.adk.auth.oauth2_credential_util import create_oauth2_session
|
||||
from google.adk.auth.oauth2_credential_util import update_credential_with_tokens
|
||||
|
||||
|
||||
class TestOAuth2CredentialUtil:
|
||||
"""Test suite for OAuth2 credential utility functions."""
|
||||
|
||||
def test_create_oauth2_session_openid_connect(self):
|
||||
"""Test create_oauth2_session with OpenID Connect scheme."""
|
||||
scheme = OpenIdConnectWithConfig(
|
||||
type_="openIdConnect",
|
||||
openId_connect_url=(
|
||||
"https://example.com/.well-known/openid_configuration"
|
||||
),
|
||||
authorization_endpoint="https://example.com/auth",
|
||||
token_endpoint="https://example.com/token",
|
||||
scopes=["openid", "profile"],
|
||||
)
|
||||
credential = AuthCredential(
|
||||
auth_type=AuthCredentialTypes.OPEN_ID_CONNECT,
|
||||
oauth2=OAuth2Auth(
|
||||
client_id="test_client_id",
|
||||
client_secret="test_client_secret",
|
||||
redirect_uri="https://example.com/callback",
|
||||
state="test_state",
|
||||
),
|
||||
)
|
||||
|
||||
client, token_endpoint = create_oauth2_session(scheme, credential)
|
||||
|
||||
assert client is not None
|
||||
assert token_endpoint == "https://example.com/token"
|
||||
assert client.client_id == "test_client_id"
|
||||
assert client.client_secret == "test_client_secret"
|
||||
|
||||
def test_create_oauth2_session_oauth2_scheme(self):
|
||||
"""Test create_oauth2_session with OAuth2 scheme."""
|
||||
flows = OAuthFlows(
|
||||
authorizationCode=OAuthFlowAuthorizationCode(
|
||||
authorizationUrl="https://example.com/auth",
|
||||
tokenUrl="https://example.com/token",
|
||||
scopes={"read": "Read access", "write": "Write access"},
|
||||
)
|
||||
)
|
||||
scheme = OAuth2(type_="oauth2", flows=flows)
|
||||
credential = AuthCredential(
|
||||
auth_type=AuthCredentialTypes.OAUTH2,
|
||||
oauth2=OAuth2Auth(
|
||||
client_id="test_client_id",
|
||||
client_secret="test_client_secret",
|
||||
redirect_uri="https://example.com/callback",
|
||||
),
|
||||
)
|
||||
|
||||
client, token_endpoint = create_oauth2_session(scheme, credential)
|
||||
|
||||
assert client is not None
|
||||
assert token_endpoint == "https://example.com/token"
|
||||
|
||||
def test_create_oauth2_session_invalid_scheme(self):
|
||||
"""Test create_oauth2_session with invalid scheme."""
|
||||
scheme = Mock() # Invalid scheme type
|
||||
credential = AuthCredential(
|
||||
auth_type=AuthCredentialTypes.OAUTH2,
|
||||
oauth2=OAuth2Auth(
|
||||
client_id="test_client_id",
|
||||
client_secret="test_client_secret",
|
||||
),
|
||||
)
|
||||
|
||||
client, token_endpoint = create_oauth2_session(scheme, credential)
|
||||
|
||||
assert client is None
|
||||
assert token_endpoint is None
|
||||
|
||||
def test_create_oauth2_session_missing_credentials(self):
|
||||
"""Test create_oauth2_session with missing credentials."""
|
||||
scheme = OpenIdConnectWithConfig(
|
||||
type_="openIdConnect",
|
||||
openId_connect_url=(
|
||||
"https://example.com/.well-known/openid_configuration"
|
||||
),
|
||||
authorization_endpoint="https://example.com/auth",
|
||||
token_endpoint="https://example.com/token",
|
||||
scopes=["openid"],
|
||||
)
|
||||
credential = AuthCredential(
|
||||
auth_type=AuthCredentialTypes.OPEN_ID_CONNECT,
|
||||
oauth2=OAuth2Auth(
|
||||
client_id="test_client_id",
|
||||
# Missing client_secret
|
||||
),
|
||||
)
|
||||
|
||||
client, token_endpoint = create_oauth2_session(scheme, credential)
|
||||
|
||||
assert client is None
|
||||
assert token_endpoint is None
|
||||
|
||||
def test_update_credential_with_tokens(self):
|
||||
"""Test update_credential_with_tokens function."""
|
||||
credential = AuthCredential(
|
||||
auth_type=AuthCredentialTypes.OPEN_ID_CONNECT,
|
||||
oauth2=OAuth2Auth(
|
||||
client_id="test_client_id",
|
||||
client_secret="test_client_secret",
|
||||
),
|
||||
)
|
||||
|
||||
tokens = OAuth2Token({
|
||||
"access_token": "new_access_token",
|
||||
"refresh_token": "new_refresh_token",
|
||||
"expires_at": int(time.time()) + 3600,
|
||||
"expires_in": 3600,
|
||||
})
|
||||
|
||||
update_credential_with_tokens(credential, tokens)
|
||||
|
||||
assert credential.oauth2.access_token == "new_access_token"
|
||||
assert credential.oauth2.refresh_token == "new_refresh_token"
|
||||
assert credential.oauth2.expires_at == int(time.time()) + 3600
|
||||
assert credential.oauth2.expires_in == 3600
|
||||
Reference in New Issue
Block a user