diff --git a/airbyte_cdk/sources/streams/http/http_client.py b/airbyte_cdk/sources/streams/http/http_client.py index caacdd503a..f79276a7aa 100644 --- a/airbyte_cdk/sources/streams/http/http_client.py +++ b/airbyte_cdk/sources/streams/http/http_client.py @@ -175,6 +175,7 @@ def __init__( self._backoff_strategies = [DefaultBackoffStrategy()] self._error_message_parser = error_message_parser or JsonErrorMessageParser() self._request_attempt_count: Dict[requests.PreparedRequest, int] = {} + self._token_refresh_outcomes: Dict[requests.PreparedRequest, bool] = {} self._disable_retries = disable_retries self._message_repository = message_repository self._authenticator_update_failed = False @@ -537,6 +538,26 @@ def _evict_key(self, prepared_request: requests.PreparedRequest) -> None: """ if prepared_request in self._request_attempt_count: del self._request_attempt_count[prepared_request] + self._token_refresh_outcomes.pop(prepared_request, None) + + def _auth_header_changed_since(self, request: requests.PreparedRequest) -> bool: + """Whether the authenticator's current Authorization header differs from the one the request was sent with.""" + # request.headers is None on an unprepared request + sent = request.headers.get("Authorization") if request.headers else None + if not sent or not hasattr(self._session.auth, "get_auth_header"): + return False + current = self._session.auth.get_auth_header().get("Authorization") # type: ignore[union-attr] + return current is not None and current != sent + + @staticmethod + def _is_token_endpoint_rejection(error: requests.exceptions.RequestException) -> bool: + """Whether the token endpoint answered with a 4xx other than 429, i.e. it rejected the credentials rather than failing transiently.""" + response = error.response + return ( + response is not None + and 400 <= response.status_code < 500 + and response.status_code != 429 + ) def _handle_error_resolution( self, @@ -575,20 +596,81 @@ def _handle_error_resolution( # backoff retry loop. Adding `\n` to the message and ignore 'end' ensure that few messages are printed at the same time. print(f"{message}\n", end="", flush=True) - # Handle REFRESH_TOKEN_THEN_RETRY: Force refresh the OAuth token before retry - # This is useful when the API returns 401 but the stored token expiry hasn't been reached yet - # Only OAuth authenticators have refresh_and_set_access_token method - # Non-OAuth auth types (e.g., BearerAuthenticator) will fall through to normal retry + # Handle REFRESH_TOKEN_THEN_RETRY: force refresh the OAuth token before retrying, + # at most once per request. A config error from the refresh (e.g. bad credentials, + # including a 4xx rejection from the token endpoint) or a request rejected again + # after a refresh attempt fails fast instead of refreshing in a loop; the failure + # type reflects whether the refresh succeeded. Transient refresh failures (network + # errors, 5xx, 429) keep the retry transient. Non-OAuth auth types (e.g., + # BearerAuthenticator) fall through to normal retry. if error_resolution.response_action == ResponseAction.REFRESH_TOKEN_THEN_RETRY: + status = ( + f"status code '{response.status_code}'" + if response is not None + else f"exception '{exc}'" + ) + if request in self._token_refresh_outcomes: + refreshed = self._token_refresh_outcomes[request] + self._evict_key(request) + if refreshed: + internal_message = f"'{request.method}' request to '{request.url}' was rejected with {status} again after the OAuth token was refreshed; not refreshing again." + failure_type = FailureType.config_error + message = "Refreshed OAuth access token is rejected by the API." + else: + internal_message = f"'{request.method}' request to '{request.url}' was rejected with {status} and the OAuth token could not be refreshed; not refreshing again." + failure_type = FailureType.transient_error + message = ( + "API rejects the current OAuth access token and the token refresh failed." + ) + if error_resolution.error_message: + internal_message += ( + f" Error handler message: '{error_resolution.error_message}'" + ) + self._logger.error(internal_message) + raise AirbyteTracedException( + internal_message=internal_message, + message=message, + failure_type=failure_type, + ) if ( hasattr(self._session, "auth") and self._session.auth is not None and hasattr(self._session.auth, "refresh_and_set_access_token") ): + self._token_refresh_outcomes[request] = False try: - self._session.auth.refresh_and_set_access_token() # type: ignore[union-attr] - self._logger.info( - "Refreshed OAuth token due to REFRESH_TOKEN_THEN_RETRY response action" + if self._auth_header_changed_since(request): + self._logger.info( + "OAuth token was already replaced since this request was sent; retrying with the current token without refreshing again." + ) + else: + self._session.auth.refresh_and_set_access_token() # type: ignore[union-attr] + self._logger.info( + "Refreshed OAuth token due to REFRESH_TOKEN_THEN_RETRY response action" + ) + self._token_refresh_outcomes[request] = True + except AirbyteTracedException as refresh_error: + if refresh_error.failure_type == FailureType.config_error: + self._evict_key(request) + raise + self._logger.warning( + f"Failed to refresh OAuth token: {refresh_error}. Proceeding with retry using existing token." + ) + except requests.exceptions.RequestException as refresh_error: + if self._is_token_endpoint_rejection(refresh_error): + self._evict_key(request) + internal_message = ( + f"'{request.method}' request to '{request.url}' was rejected with {status} and the OAuth token endpoint " + f"rejected the refresh request: {refresh_error}" + ) + self._logger.error(internal_message) + raise AirbyteTracedException( + internal_message=internal_message, + message="OAuth token refresh request is rejected by the token endpoint.", + failure_type=FailureType.config_error, + ) from refresh_error + self._logger.warning( + f"Failed to refresh OAuth token: {refresh_error}. Proceeding with retry using existing token." ) except Exception as refresh_error: self._logger.warning( diff --git a/airbyte_cdk/sources/streams/http/requests_native_auth/abstract_oauth.py b/airbyte_cdk/sources/streams/http/requests_native_auth/abstract_oauth.py index e236463814..ff842c9ce4 100644 --- a/airbyte_cdk/sources/streams/http/requests_native_auth/abstract_oauth.py +++ b/airbyte_cdk/sources/streams/http/requests_native_auth/abstract_oauth.py @@ -61,7 +61,7 @@ class AbstractOauth2Authenticator(AuthBase): # through the connector config. Without this lock, concurrent refresh attempts can cause race # conditions where one stream successfully refreshes the token while others fail because the # refresh token has been invalidated (especially for single-use refresh tokens). - _token_refresh_lock: threading.Lock = threading.Lock() + _token_refresh_lock = threading.RLock() def __init__( self, @@ -124,16 +124,36 @@ def get_access_token(self) -> str: return self.access_token + def _current_access_token_or_none(self) -> Optional[str]: + """The current access token, or None when the implementation has none to report + (e.g., a declarative authenticator whose token has not been initialized yet).""" + try: + return self.access_token + except Exception: + return None + def refresh_and_set_access_token(self) -> None: """Force refresh the access token and update internal state. - This method refreshes the access token regardless of whether it has expired, - and updates the internal token and expiry date. Subclasses may override this - to handle additional state updates (e.g., persisting new refresh tokens). - """ - token, expires_in = self.refresh_access_token() - self.access_token = token - self.set_token_expiry_date(expires_in) + Refreshes regardless of expiry, serialized on the class-level refresh lock. If another + thread using this same authenticator instance replaced the access token while this one + waited for the lock, the refresh is skipped and the request is retried with that token. + `SingleUseRefreshTokenOauth2Authenticator` reads `access_token` from the connector config + shared by all stream instances, so the early return also covers separate instances there. + Only per-instance token authenticators (base and declarative) are limited to same-instance + detection; the Authorization-header check in `HttpClient._handle_error_resolution` covers + the rest by skipping the forced refresh when the rejected request's token was already + replaced. + Subclasses may override this to handle additional state updates (e.g., persisting new + refresh tokens). + """ + token_before_waiting = self._current_access_token_or_none() + with self._token_refresh_lock: + if self._current_access_token_or_none() != token_before_waiting: + return + token, expires_in = self.refresh_access_token() + self.access_token = token + self.set_token_expiry_date(expires_in) def token_has_expired(self) -> bool: """Returns True if the token is expired""" diff --git a/airbyte_cdk/sources/streams/http/requests_native_auth/oauth.py b/airbyte_cdk/sources/streams/http/requests_native_auth/oauth.py index e4243b3a28..59e46bdfac 100644 --- a/airbyte_cdk/sources/streams/http/requests_native_auth/oauth.py +++ b/airbyte_cdk/sources/streams/http/requests_native_auth/oauth.py @@ -351,12 +351,18 @@ def refresh_and_set_access_token(self) -> None: and emits a control message to update the connector config. If the response omits a refresh token, the existing one is preserved. """ - new_access_token, access_token_expires_in, new_refresh_token = self.refresh_access_token() - self.access_token = new_access_token - if new_refresh_token is not None: - self.set_refresh_token(new_refresh_token) - self.set_token_expiry_date(access_token_expires_in) - self._emit_control_message() + token_before_waiting = self._current_access_token_or_none() + with self._token_refresh_lock: + if self._current_access_token_or_none() != token_before_waiting: + return + new_access_token, access_token_expires_in, new_refresh_token = ( + self.refresh_access_token() + ) + self.access_token = new_access_token + if new_refresh_token is not None: + self.set_refresh_token(new_refresh_token) + self.set_token_expiry_date(access_token_expires_in) + self._emit_control_message() def refresh_access_token(self) -> Tuple[str, AirbyteDateTime, Optional[str]]: # type: ignore[override] """ diff --git a/unit_tests/sources/streams/http/requests_native_auth/test_requests_native_auth.py b/unit_tests/sources/streams/http/requests_native_auth/test_requests_native_auth.py index 98450e7f01..bbc6a17086 100644 --- a/unit_tests/sources/streams/http/requests_native_auth/test_requests_native_auth.py +++ b/unit_tests/sources/streams/http/requests_native_auth/test_requests_native_auth.py @@ -1414,3 +1414,197 @@ def get_token(): assert len(results) == 5 assert all(token == "new_access_token" for token in results) assert refresh_call_count == 1, f"Expected 1 refresh call, got {refresh_call_count}" + + def test_refresh_and_set_access_token_skips_refresh_when_another_thread_already_refreshed( + self, mocker + ): + """ + A forced refresh that was queued behind the class-level lock must not refresh + again when the access token changed while it waited: the request is retried + with the token the other thread already obtained. + """ + oauth = Oauth2Authenticator( + token_refresh_endpoint="https://refresh_endpoint.com", + client_id="client_id", + client_secret="client_secret", + refresh_token="refresh_token", + token_expiry_date=ab_datetime_now() + timedelta(hours=1), + ) + oauth.access_token = "old" + + mocked_refresh = mocker.patch.object(oauth, "refresh_access_token") + + snapshot_taken = threading.Event() + real_snapshot = oauth._current_access_token_or_none + + def record_snapshot(): + value = real_snapshot() + snapshot_taken.set() + return value + + mocker.patch.object(oauth, "_current_access_token_or_none", side_effect=record_snapshot) + + with Oauth2Authenticator._token_refresh_lock: + thread = threading.Thread(target=oauth.refresh_and_set_access_token) + thread.start() + # The worker took its pre-lock token snapshot and is now blocked on the lock. + assert snapshot_taken.wait(5) + oauth.access_token = "winner" + thread.join(timeout=5) + + assert not thread.is_alive() + mocked_refresh.assert_not_called() + assert oauth.access_token == "winner" + + def test_refresh_and_set_access_token_refreshes_when_token_unchanged(self, mocker): + oauth = Oauth2Authenticator( + token_refresh_endpoint="https://refresh_endpoint.com", + client_id="client_id", + client_secret="client_secret", + refresh_token="refresh_token", + token_expiry_date=ab_datetime_now() + timedelta(hours=1), + ) + oauth.access_token = "old" + + mocked_refresh = mocker.patch.object( + oauth, + "refresh_access_token", + return_value=("new", ab_datetime_now() + timedelta(hours=1)), + ) + + oauth.refresh_and_set_access_token() + + mocked_refresh.assert_called_once() + assert oauth.access_token == "new" + + def test_get_access_token_refresh_lock_is_reentrant(self, requests_mock): + """ + get_access_token() holds the class-level refresh lock while calling + refresh_and_set_access_token, which re-acquires it; a non-reentrant lock + would deadlock. The worker thread must finish and report the new token. + """ + requests_mock.post( + "https://refresh_endpoint.com", + json={"access_token": "new_access_token", "expires_in": 3600}, + ) + oauth = Oauth2Authenticator( + token_refresh_endpoint="https://refresh_endpoint.com", + client_id="client_id", + client_secret="client_secret", + refresh_token="refresh_token", + token_expiry_date=ab_datetime_now() - timedelta(hours=1), + ) + + results, errors = [], [] + + def get_token(): + try: + results.append(oauth.get_access_token()) + except Exception as e: + errors.append(e) + + thread = threading.Thread(target=get_token) + thread.start() + thread.join(timeout=5) + + assert not thread.is_alive(), "get_access_token deadlocked on a non-reentrant lock" + assert errors == [] + assert results == ["new_access_token"] + + def test_refresh_and_set_access_token_holds_lock_during_refresh(self, mocker): + """ + The class-level lock must be held for the whole refresh: a probe from another + thread cannot acquire it while refresh_access_token is in flight. + """ + oauth = Oauth2Authenticator( + token_refresh_endpoint="https://refresh_endpoint.com", + client_id="client_id", + client_secret="client_secret", + refresh_token="refresh_token", + token_expiry_date=ab_datetime_now() + timedelta(hours=1), + ) + oauth.access_token = "old" + + probe_results = [] + + def spy_refresh_access_token(self): + def probe(): + acquired = Oauth2Authenticator._token_refresh_lock.acquire(blocking=False) + if acquired: + Oauth2Authenticator._token_refresh_lock.release() + probe_results.append(acquired) + + probe_thread = threading.Thread(target=probe) + probe_thread.start() + probe_thread.join() + return ("new", ab_datetime_now() + timedelta(hours=1)) + + mocker.patch.object(Oauth2Authenticator, "refresh_access_token", spy_refresh_access_token) + + oauth.refresh_and_set_access_token() + + assert probe_results == [False] + assert oauth.access_token == "new" + + +class TestSingleUseRefreshTokenInstanceSkip: + """ + Two SingleUseRefreshTokenOauth2Authenticator instances over one shared connector + config: a forced refresh queued behind the lock sees the token the first instance + already wrote back into the shared config and skips its own refresh. + """ + + def test_second_instance_skips_refresh_after_first_refreshes(self, requests_mock, mocker): + requests_mock.post( + "https://refresh_endpoint.com", + json={ + "access_token": "new_access_token", + "refresh_token": "new_refresh_token", + "expires_in": 3600, + }, + ) + connector_config = { + "credentials": { + "client_id": "client_id", + "client_secret": "client_secret", + "refresh_token": "refresh_token", + "access_token": "old_access_token", + "token_expiry_date": str(ab_datetime_now() + timedelta(hours=1)), + } + } + auth1 = SingleUseRefreshTokenOauth2Authenticator( + connector_config=connector_config, + token_refresh_endpoint="https://refresh_endpoint.com", + ) + auth2 = SingleUseRefreshTokenOauth2Authenticator( + connector_config=connector_config, + token_refresh_endpoint="https://refresh_endpoint.com", + ) + + mocked_emit = mocker.patch.object( + SingleUseRefreshTokenOauth2Authenticator, "_emit_control_message" + ) + + snapshot_taken = threading.Event() + real_snapshot = auth2._current_access_token_or_none + + def record_snapshot(): + value = real_snapshot() + snapshot_taken.set() + return value + + mocker.patch.object(auth2, "_current_access_token_or_none", side_effect=record_snapshot) + + with SingleUseRefreshTokenOauth2Authenticator._token_refresh_lock: + thread = threading.Thread(target=auth2.refresh_and_set_access_token) + thread.start() + # auth2 took its pre-lock snapshot ("old_access_token") and now waits on the lock. + assert snapshot_taken.wait(5) + auth1.refresh_and_set_access_token() + thread.join(timeout=5) + + assert not thread.is_alive() + assert len(requests_mock.request_history) == 1 + assert mocked_emit.call_count == 1 + assert auth2.access_token == "new_access_token" + assert connector_config["credentials"]["refresh_token"] == "new_refresh_token" diff --git a/unit_tests/sources/streams/http/test_http_client.py b/unit_tests/sources/streams/http/test_http_client.py index 14ddcaf7a6..8724c48e49 100644 --- a/unit_tests/sources/streams/http/test_http_client.py +++ b/unit_tests/sources/streams/http/test_http_client.py @@ -13,6 +13,7 @@ from requests_cache import CachedRequest from airbyte_cdk.models import FailureType, Level +from airbyte_cdk.sources.declarative.auth.oauth import DeclarativeOauth2Authenticator from airbyte_cdk.sources.http_logger import format_http_message from airbyte_cdk.sources.message import InMemoryMessageRepository from airbyte_cdk.sources.streams.call_rate import CachedLimiterSession, LimiterSession @@ -33,7 +34,11 @@ from airbyte_cdk.sources.streams.http.page_size_reduction_exception import ( PageSizeReductionRequiredException, ) -from airbyte_cdk.sources.streams.http.requests_native_auth import TokenAuthenticator +from airbyte_cdk.sources.streams.http.requests_native_auth import ( + Oauth2Authenticator, + TokenAuthenticator, +) +from airbyte_cdk.utils.datetime_helpers import ab_datetime_now from airbyte_cdk.utils.traced_exception import AirbyteTracedException @@ -1069,6 +1074,274 @@ def update_response(*args, **kwargs): assert call_count == 2 +def _build_refresh_token_then_retry_http_client(authenticator=None): + """An HttpClient backed by a real Oauth2Authenticator with a non-expired token, so the + first request does not itself trigger a refresh and only REFRESH_TOKEN_THEN_RETRY does.""" + if authenticator is None: + authenticator = Oauth2Authenticator( + token_refresh_endpoint="https://example.com/oauth/token", + client_id="client_id", + client_secret="client_secret", + refresh_token="refresh_token", + token_expiry_date=ab_datetime_now() + timedelta(days=1), + refresh_token_error_status_codes=(400,), + refresh_token_error_key="error", + refresh_token_error_values=("invalid_grant",), + ) + http_client = HttpClient( + name="test", + logger=logging.getLogger("test"), + authenticator=authenticator, + error_handler=HttpStatusErrorHandler( + logger=logging.getLogger("test"), + max_retries=5, + error_mapping={ + 401: ErrorResolution( + ResponseAction.REFRESH_TOKEN_THEN_RETRY, + FailureType.transient_error, + "Token rejected; refresh and retry", + ) + }, + ), + ) + return http_client + + +def _request_count(requests_mock, url, method="GET"): + return len([r for r in requests_mock.request_history if r.url == url and r.method == method]) + + +@pytest.mark.usefixtures("mock_sleep") +def test_refresh_token_then_retry_refreshes_once_and_succeeds(requests_mock): + requests_mock.get( + "https://example.com/data", + [{"status_code": 401}, {"status_code": 200, "json": {"data": "ok"}}], + ) + requests_mock.post( + "https://example.com/oauth/token", + json={"access_token": "new", "expires_in": 3600}, + ) + http_client = _build_refresh_token_then_retry_http_client() + + _, response = http_client.send_request("GET", "https://example.com/data", request_kwargs={}) + + assert response.status_code == 200 + assert _request_count(requests_mock, "https://example.com/oauth/token", "POST") == 1 + assert _request_count(requests_mock, "https://example.com/data") == 2 + second_request = [ + r for r in requests_mock.request_history if r.url == "https://example.com/data" + ][1] + assert second_request.headers["Authorization"] == "Bearer new" + + +def test_refresh_token_then_retry_fails_fast_when_refresh_is_rejected(requests_mock): + requests_mock.get("https://example.com/data", status_code=401) + requests_mock.post( + "https://example.com/oauth/token", + status_code=400, + json={"error": "invalid_grant"}, + ) + http_client = _build_refresh_token_then_retry_http_client() + + with patch("time.sleep") as mocked_sleep: + with pytest.raises(AirbyteTracedException) as exc_info: + http_client.send_request("GET", "https://example.com/data", request_kwargs={}) + + assert exc_info.value.failure_type == FailureType.config_error + assert _request_count(requests_mock, "https://example.com/oauth/token", "POST") == 1 + assert _request_count(requests_mock, "https://example.com/data") == 1 + assert http_client._token_refresh_outcomes == {} + mocked_sleep.assert_not_called() + + +@pytest.mark.usefixtures("mock_sleep") +def test_refresh_token_then_retry_does_not_refresh_twice_for_the_same_request(requests_mock): + requests_mock.get("https://example.com/data", status_code=401) + requests_mock.post( + "https://example.com/oauth/token", + json={"access_token": "new", "expires_in": 3600}, + ) + http_client = _build_refresh_token_then_retry_http_client() + + with pytest.raises(AirbyteTracedException) as exc_info: + http_client.send_request("GET", "https://example.com/data", request_kwargs={}) + + assert exc_info.value.failure_type == FailureType.config_error + assert exc_info.value.message == "Refreshed OAuth access token is rejected by the API." + assert "Token rejected; refresh and retry" in exc_info.value.internal_message + assert _request_count(requests_mock, "https://example.com/oauth/token", "POST") == 1 + assert _request_count(requests_mock, "https://example.com/data") == 2 + assert http_client._token_refresh_outcomes == {} + + +def test_refresh_token_then_retry_reports_transient_error_when_refresh_fails_transiently( + requests_mock, mocker +): + requests_mock.get("https://example.com/data", status_code=401) + requests_mock.post( + "https://example.com/oauth/token", + exc=requests.exceptions.ConnectionError("token endpoint unreachable"), + ) + http_client = _build_refresh_token_then_retry_http_client() + + mocker.patch("time.sleep") + with pytest.raises(AirbyteTracedException) as exc_info: + http_client.send_request("GET", "https://example.com/data", request_kwargs={}) + + assert exc_info.value.failure_type == FailureType.transient_error + assert ( + exc_info.value.message + == "API rejects the current OAuth access token and the token refresh failed." + ) + # The token endpoint itself is retried by the backoff decorator on _make_handled_request. + assert _request_count(requests_mock, "https://example.com/oauth/token", "POST") > 1 + assert _request_count(requests_mock, "https://example.com/data") == 2 + + +@pytest.mark.usefixtures("mock_sleep") +def test_refresh_token_then_retry_refresh_rejected_outside_configured_errors(requests_mock): + requests_mock.get("https://example.com/data", status_code=401) + requests_mock.post( + "https://example.com/oauth/token", + status_code=401, + json={"error": "invalid_client"}, + ) + http_client = _build_refresh_token_then_retry_http_client() + + with pytest.raises(AirbyteTracedException) as exc_info: + http_client.send_request("GET", "https://example.com/data", request_kwargs={}) + + assert exc_info.value.failure_type == FailureType.config_error + assert ( + exc_info.value.message == "OAuth token refresh request is rejected by the token endpoint." + ) + assert "401" in exc_info.value.internal_message + assert _request_count(requests_mock, "https://example.com/oauth/token", "POST") == 1 + assert _request_count(requests_mock, "https://example.com/data") == 1 + assert http_client._token_refresh_outcomes == {} + + +def test_refresh_token_then_retry_token_endpoint_5xx_stays_transient(requests_mock, mocker): + """A 5xx from the token endpoint is a transient refresh failure, not a credential + rejection: the refresh is not retried here (the endpoint's own backoff decorator + handles that) and the request stays on the warn-and-retry path.""" + requests_mock.get("https://example.com/data", status_code=401) + http_client = _build_refresh_token_then_retry_http_client() + + token_response = requests.Response() + token_response.status_code = 503 + token_response.url = "https://example.com/oauth/token" + refresh_error = DefaultBackoffException( + request=requests.Request(method="POST", url="https://example.com/oauth/token").prepare(), + response=token_response, + failure_type=FailureType.transient_error, + ) + mocker.patch.object(Oauth2Authenticator, "_make_handled_request", side_effect=refresh_error) + mocker.patch("time.sleep") + + with pytest.raises(AirbyteTracedException) as exc_info: + http_client.send_request("GET", "https://example.com/data", request_kwargs={}) + + assert exc_info.value.failure_type == FailureType.transient_error + assert ( + exc_info.value.message + == "API rejects the current OAuth access token and the token refresh failed." + ) + assert _request_count(requests_mock, "https://example.com/data") == 2 + + +@pytest.mark.usefixtures("mock_sleep") +def test_refresh_token_then_retry_skips_refresh_when_token_already_replaced(requests_mock): + http_client = _build_refresh_token_then_retry_http_client() + authenticator = http_client._session.auth + + def expire_token(request, context): + authenticator.access_token = "new" + context.status_code = 401 + return json.dumps({}) + + requests_mock.get( + "https://example.com/data", + [{"text": expire_token}, {"status_code": 200, "json": {"data": "ok"}}], + ) + + with patch.object(authenticator, "refresh_and_set_access_token") as refresh_spy: + _, response = http_client.send_request("GET", "https://example.com/data", request_kwargs={}) + + assert response.status_code == 200 + refresh_spy.assert_not_called() + assert _request_count(requests_mock, "https://example.com/oauth/token", "POST") == 0 + assert _request_count(requests_mock, "https://example.com/data") == 2 + second_request = [ + r for r in requests_mock.request_history if r.url == "https://example.com/data" + ][1] + assert second_request.headers["Authorization"] == "Bearer new" + assert http_client._token_refresh_outcomes == {} + + +@pytest.mark.usefixtures("mock_sleep") +def test_refresh_token_then_retry_with_declarative_oauth_authenticator(requests_mock): + authenticator = DeclarativeOauth2Authenticator( + token_refresh_endpoint="https://example.com/oauth/token", + client_id="client_id", + client_secret="client_secret", + refresh_token="refresh_token", + config={}, + parameters={}, + access_token_value="old", + token_expiry_date=(ab_datetime_now() + timedelta(days=1)).isoformat(), + refresh_token_error_status_codes=(400,), + refresh_token_error_key="error", + refresh_token_error_values=("invalid_grant",), + ) + http_client = _build_refresh_token_then_retry_http_client(authenticator=authenticator) + + requests_mock.get( + "https://example.com/data", + [{"status_code": 401}, {"status_code": 200, "json": {"data": "ok"}}], + ) + requests_mock.post( + "https://example.com/oauth/token", + json={"access_token": "new", "expires_in": 3600}, + ) + + _, response = http_client.send_request("GET", "https://example.com/data", request_kwargs={}) + + assert response.status_code == 200 + assert response.json() == {"data": "ok"} + assert _request_count(requests_mock, "https://example.com/oauth/token", "POST") == 1 + assert _request_count(requests_mock, "https://example.com/data") == 2 + + +@pytest.mark.usefixtures("mock_sleep") +def test_refresh_token_then_retry_state_is_evicted_after_success(requests_mock): + requests_mock.get( + "https://example.com/data", + [{"status_code": 401}, {"status_code": 200, "json": {"data": "ok"}}], + ) + requests_mock.get( + "https://example.com/other", + [{"status_code": 401}, {"status_code": 200, "json": {"data": "ok"}}], + ) + requests_mock.post( + "https://example.com/oauth/token", + json={"access_token": "new", "expires_in": 3600}, + ) + http_client = _build_refresh_token_then_retry_http_client() + + _, response = http_client.send_request("GET", "https://example.com/data", request_kwargs={}) + assert response.status_code == 200 + assert http_client._token_refresh_outcomes == {} + # A different URL for the second request: PreparedRequest instances are keyed by + # identity, but a distinct URL also keeps the two request histories easy to count. + _, second_response = http_client.send_request( + "GET", "https://example.com/other", request_kwargs={} + ) + assert second_response.status_code == 200 + + assert _request_count(requests_mock, "https://example.com/oauth/token", "POST") == 2 + + class _RecordingAuthenticator(TokenAuthenticator): """An authenticator that tracks quota state and wants to see responses."""