|
16 | 16 | from unittest.mock import Mock |
17 | 17 | from unittest.mock import patch |
18 | 18 |
|
| 19 | +from authlib.integrations.base_client.errors import OAuthError |
19 | 20 | from authlib.oauth2.rfc6749 import OAuth2Token |
| 21 | +from authlib.oauth2.rfc6749.errors import InvalidGrantError |
20 | 22 | from authlib.oauth2.rfc6749.errors import OAuth2Error |
21 | 23 | from google.adk.auth.auth_credential import AuthCredential |
22 | 24 | from google.adk.auth.auth_credential import AuthCredentialTypes |
23 | 25 | from google.adk.auth.auth_credential import OAuth2Auth |
24 | 26 | from google.adk.auth.auth_schemes import OpenIdConnectWithConfig |
| 27 | +from google.adk.auth.refresher import base_credential_refresher |
25 | 28 | from google.adk.auth.refresher.oauth2_credential_refresher import OAuth2CredentialRefresher |
26 | 29 | import pytest |
27 | 30 | import requests |
@@ -195,7 +198,7 @@ async def test_refresh_oauth2_error_returns_original_and_logs( |
195 | 198 | mock_client = Mock() |
196 | 199 | mock_oauth2_session.return_value = mock_client |
197 | 200 | mock_client.refresh_token.side_effect = OAuth2Error( |
198 | | - description="invalid_grant" |
| 201 | + description="temporarily_unavailable" |
199 | 202 | ) |
200 | 203 |
|
201 | 204 | scheme = OpenIdConnectWithConfig( |
@@ -309,6 +312,94 @@ async def test_refresh_unexpected_error_propagates( |
309 | 312 | with pytest.raises(ValueError, match="unexpected bug"): |
310 | 313 | await refresher.refresh(credential, scheme) |
311 | 314 |
|
| 315 | + @patch("google.adk.auth.oauth2_credential_util.OAuth2Session") |
| 316 | + @patch("google.adk.auth.oauth2_credential_util.OAuth2Token") |
| 317 | + @pytest.mark.parametrize( |
| 318 | + "error", |
| 319 | + [ |
| 320 | + InvalidGrantError(description="Token has been revoked."), |
| 321 | + OAuth2Error(error="invalid_grant"), |
| 322 | + OAuthError(error="invalid_grant", description="Bad Request"), |
| 323 | + ], |
| 324 | + ) |
| 325 | + @pytest.mark.asyncio |
| 326 | + async def test_refresh_invalid_grant_raises_revoked( |
| 327 | + self, mock_oauth2_token, mock_oauth2_session, error |
| 328 | + ): |
| 329 | + """invalid_grant means the refresh token is dead and must not be kept.""" |
| 330 | + mock_token_instance = Mock() |
| 331 | + mock_token_instance.is_expired.return_value = True |
| 332 | + mock_oauth2_token.return_value = mock_token_instance |
| 333 | + |
| 334 | + mock_client = Mock() |
| 335 | + mock_oauth2_session.return_value = mock_client |
| 336 | + mock_client.refresh_token.side_effect = error |
| 337 | + |
| 338 | + scheme = OpenIdConnectWithConfig( |
| 339 | + type_="openIdConnect", |
| 340 | + openId_connect_url=( |
| 341 | + "https://example.com/.well-known/openid_configuration" |
| 342 | + ), |
| 343 | + authorization_endpoint="https://example.com/auth", |
| 344 | + token_endpoint="https://example.com/token", |
| 345 | + scopes=["openid"], |
| 346 | + ) |
| 347 | + credential = AuthCredential( |
| 348 | + auth_type=AuthCredentialTypes.OPEN_ID_CONNECT, |
| 349 | + oauth2=OAuth2Auth( |
| 350 | + client_id="test_client_id", |
| 351 | + client_secret="test_client_secret", |
| 352 | + access_token="old_token", |
| 353 | + refresh_token="old_refresh_token", |
| 354 | + expires_at=int(time.time()) - 3600, # Expired |
| 355 | + ), |
| 356 | + ) |
| 357 | + |
| 358 | + refresher = OAuth2CredentialRefresher() |
| 359 | + with pytest.raises(base_credential_refresher._CredentialRevokedError): |
| 360 | + await refresher.refresh(credential, scheme) |
| 361 | + |
| 362 | + @patch("google.adk.auth.oauth2_credential_util.OAuth2Session") |
| 363 | + @patch("google.adk.auth.oauth2_credential_util.OAuth2Token") |
| 364 | + @pytest.mark.asyncio |
| 365 | + async def test_refresh_other_oauth2_error_returns_original( |
| 366 | + self, mock_oauth2_token, mock_oauth2_session |
| 367 | + ): |
| 368 | + """An OAuth2 error other than invalid_grant is transient: keep the token.""" |
| 369 | + mock_token_instance = Mock() |
| 370 | + mock_token_instance.is_expired.return_value = True |
| 371 | + mock_oauth2_token.return_value = mock_token_instance |
| 372 | + |
| 373 | + mock_client = Mock() |
| 374 | + mock_oauth2_session.return_value = mock_client |
| 375 | + mock_client.refresh_token.side_effect = OAuth2Error(error="server_error") |
| 376 | + |
| 377 | + scheme = OpenIdConnectWithConfig( |
| 378 | + type_="openIdConnect", |
| 379 | + openId_connect_url=( |
| 380 | + "https://example.com/.well-known/openid_configuration" |
| 381 | + ), |
| 382 | + authorization_endpoint="https://example.com/auth", |
| 383 | + token_endpoint="https://example.com/token", |
| 384 | + scopes=["openid"], |
| 385 | + ) |
| 386 | + credential = AuthCredential( |
| 387 | + auth_type=AuthCredentialTypes.OPEN_ID_CONNECT, |
| 388 | + oauth2=OAuth2Auth( |
| 389 | + client_id="test_client_id", |
| 390 | + client_secret="test_client_secret", |
| 391 | + access_token="old_token", |
| 392 | + refresh_token="old_refresh_token", |
| 393 | + expires_at=int(time.time()) - 3600, # Expired |
| 394 | + ), |
| 395 | + ) |
| 396 | + |
| 397 | + refresher = OAuth2CredentialRefresher() |
| 398 | + result = await refresher.refresh(credential, scheme) |
| 399 | + |
| 400 | + assert result is credential |
| 401 | + assert result.oauth2.access_token == "old_token" |
| 402 | + |
312 | 403 | @patch("google.adk.auth.refresher.oauth2_credential_refresher.logger") |
313 | 404 | @pytest.mark.asyncio |
314 | 405 | async def test_refresh_without_refresh_token_returns_original_and_logs( |
|
0 commit comments