diff --git a/packages/google-auth/google/auth/_regional_access_boundary_utils.py b/packages/google-auth/google/auth/_regional_access_boundary_utils.py index 4e25bcc99412..3c56f4696e1a 100644 --- a/packages/google-auth/google/auth/_regional_access_boundary_utils.py +++ b/packages/google-auth/google/auth/_regional_access_boundary_utils.py @@ -97,6 +97,7 @@ def __init__(self): ) self.refresh_manager = _RegionalAccessBoundaryRefreshManager() self._update_lock = threading.Lock() + self._use_blocking_regional_access_boundary_lookup = False def __getstate__(self): """Pickle helper that serializes the _update_lock attribute.""" @@ -109,6 +110,36 @@ def __setstate__(self, state): self.__dict__.update(state) self._update_lock = threading.Lock() + def __eq__(self, other): + """Checks if two managers are equal.""" + if not isinstance(other, _RegionalAccessBoundaryManager): + return NotImplemented + return ( + self._data == other._data + and self._use_blocking_regional_access_boundary_lookup + == other._use_blocking_regional_access_boundary_lookup + ) + + def use_blocking_regional_access_boundary_lookup(self): + """Enables blocking regional access boundary lookup to true""" + self._use_blocking_regional_access_boundary_lookup = True + + def set_initial_regional_access_boundary(self, seed): + """Manually sets the regional access boundary to the client provided seed + + Args: + seed (Mapping[str, str]): The regional access boundary to use for the + credential. This should be a map with, at a minimum, an "encodedLocations" + key that maps to a hex string and an "expiry" key which maps to a + datetime.datetime. + """ + self._data = _RegionalAccessBoundaryData( + encoded_locations=seed.get("encodedLocations", None), + expiry=seed.get("expiry", None), + cooldown_expiry=None, + cooldown_duration=DEFAULT_REGIONAL_ACCESS_BOUNDARY_COOLDOWN, + ) + def apply_headers(self, headers): """Applies the Regional Access Boundary header to the provided dictionary. @@ -151,48 +182,47 @@ def maybe_start_refresh(self, credentials, request): return # If all checks pass, start the background refresh. - self.refresh_manager.start_refresh(credentials, request, self) - - -class _RegionalAccessBoundaryRefreshThread(threading.Thread): - """Thread for background refreshing of the Regional Access Boundary.""" - - def __init__(self, credentials, request, rab_manager): - super().__init__() - self.daemon = True - self._credentials = credentials - self._request = request - self._rab_manager = rab_manager + if self._use_blocking_regional_access_boundary_lookup: + self.start_blocking_refresh(credentials, request) + else: + self.refresh_manager.start_refresh(credentials, request, self) - def run(self): - """ - Performs the Regional Access Boundary lookup and updates the state. + def start_blocking_refresh(self, credentials, request): + """Initiates a blocking lookup of the Regional Access Boundary. - This method is run in a separate thread. It delegates the actual lookup - to the credentials object's `_lookup_regional_access_boundary` method. - Based on the lookup's outcome (success or complete failure after retries), - it updates the cached Regional Access Boundary information, - its expiry, its cooldown expiry, and its exponential cooldown duration. + Args: + credentials (google.auth.credentials.Credentials): The credentials to refresh. + request (google.auth.transport.Request): The object used to make HTTP requests. """ - # Catch exceptions (e.g., from the underlying transport) to prevent the - # background thread from crashing. This ensures we can gracefully enter - # an exponential cooldown state on failure. try: + # A blocking parameter is passed here to indicate this is a blocking lookup, + # which in turn will do two things: 1) set a timeout to 3s instead of the + # default 120s and 2) ensure we do not retry at all + blocking = True regional_access_boundary_info = ( - self._credentials._lookup_regional_access_boundary(self._request) + credentials._lookup_regional_access_boundary(request, blocking) ) except Exception as e: if _helpers.is_logging_enabled(_LOGGER): _LOGGER.warning( - "Asynchronous Regional Access Boundary lookup raised an exception: %s", + "Blocking Regional Access Boundary lookup raised an exception: %s", e, exc_info=True, ) regional_access_boundary_info = None - with self._rab_manager._update_lock: + self.process_regional_access_boundary_info(regional_access_boundary_info) + + def process_regional_access_boundary_info(self, regional_access_boundary_info): + """Processes the regional access boundary info and updates the state. + + Args: + regional_access_boundary_info (Optional[Mapping[str, str]]): The regional access + boundary info to process. + """ + with self._update_lock: # Capture the current state before calculating updates. - current_data = self._rab_manager._data + current_data = self._data if regional_access_boundary_info: # On success, update the boundary and its expiry, and clear any cooldown. @@ -206,14 +236,12 @@ def run(self): cooldown_duration=DEFAULT_REGIONAL_ACCESS_BOUNDARY_COOLDOWN, ) if _helpers.is_logging_enabled(_LOGGER): - _LOGGER.debug( - "Asynchronous Regional Access Boundary lookup successful." - ) + _LOGGER.debug("Regional Access Boundary lookup successful.") else: # On failure, calculate cooldown and update state. if _helpers.is_logging_enabled(_LOGGER): _LOGGER.warning( - "Asynchronous Regional Access Boundary lookup failed. Entering cooldown." + "Regional Access Boundary lookup failed. Entering cooldown." ) next_cooldown_expiry = ( @@ -241,7 +269,48 @@ def run(self): ) # Perform the atomic swap of the state object. - self._rab_manager._data = updated_data + self._data = updated_data + + +class _RegionalAccessBoundaryRefreshThread(threading.Thread): + """Thread for background refreshing of the Regional Access Boundary.""" + + def __init__(self, credentials, request, rab_manager): + super().__init__() + self.daemon = True + self._credentials = credentials + self._request = request + self._rab_manager = rab_manager + + def run(self): + """ + Performs the Regional Access Boundary lookup and updates the state. + + This method is run in a separate thread. It delegates the actual lookup + to the credentials object's `_lookup_regional_access_boundary` method. + Based on the lookup's outcome (success or complete failure after retries), + it updates the cached Regional Access Boundary information, + its expiry, its cooldown expiry, and its exponential cooldown duration. + """ + # Catch exceptions (e.g., from the underlying transport) to prevent the + # background thread from crashing. This ensures we can gracefully enter + # an exponential cooldown state on failure. + try: + regional_access_boundary_info = ( + self._credentials._lookup_regional_access_boundary(self._request) + ) + except Exception as e: + if _helpers.is_logging_enabled(_LOGGER): + _LOGGER.warning( + "Asynchronous Regional Access Boundary lookup raised an exception: %s", + e, + exc_info=True, + ) + regional_access_boundary_info = None + + self._rab_manager.process_regional_access_boundary_info( + regional_access_boundary_info + ) class _RegionalAccessBoundaryRefreshManager(object): diff --git a/packages/google-auth/google/auth/credentials.py b/packages/google-auth/google/auth/credentials.py index 1e16ca2e87a7..a277ff435e1e 100644 --- a/packages/google-auth/google/auth/credentials.py +++ b/packages/google-auth/google/auth/credentials.py @@ -361,6 +361,32 @@ def _copy_regional_access_boundary_manager(self, target): new_manager._data = self._rab_manager._data target._rab_manager = new_manager + def _with_regional_access_boundary(self, seed): + """Returns a copy of these credentials with the the regional_access_boundary + set to the provided seed. This is intended for internal use only as invalid + seeds would produce unexpected results until automatic recovery is supported. + Currently this is used by the gcloud CLI and therefore changes to the + contract MUST be backwards compatible (e.g. the method signature must be + unchanged and a copy of the credenials with the RAB set must be returned). + + + Returns: + google.auth.credentials.Credentials: A new credentials instance. + """ + creds = self._make_copy() + creds._rab_manager.set_initial_regional_access_boundary(seed) + return creds + + def with_blocking_regional_access_boundary_lookup(self): + """Returns a copy of these credentials with the blocking lookup mode enabled. + + Returns: + google.auth.credentials.Credentials: A new credentials instance. + """ + creds = self._make_copy() + creds._rab_manager.use_blocking_regional_access_boundary_lookup() + return creds + def _maybe_start_regional_access_boundary_refresh(self, request, url): """ Starts a background thread to refresh the Regional Access Boundary if needed. @@ -421,11 +447,16 @@ def before_request(self, request, method, url, headers): """Refreshes the access token and triggers the Regional Access Boundary lookup if necessary. """ - super(CredentialsWithRegionalAccessBoundary, self).before_request( - request, method, url, headers - ) + if self._use_non_blocking_refresh: + self._non_blocking_refresh(request) + else: + self._blocking_refresh(request) + self._maybe_start_regional_access_boundary_refresh(request, url) + metrics.add_metric_header(headers, self._metric_header_for_usage()) + self.apply(headers) + def refresh(self, request): """Refreshes the access token. @@ -435,13 +466,16 @@ def refresh(self, request): self._perform_refresh_token(request) def _lookup_regional_access_boundary( - self, request: "google.auth.transport.Request" # noqa: F821 + self, + request: "google.auth.transport.Request", # noqa: F821 + blocking: bool = False, ) -> "Optional[Dict[str, str]]": """Calls the Regional Access Boundary lookup API to retrieve the Regional Access Boundary information. Args: request (google.auth.transport.Request): The object used to make HTTP requests. + blocking (bool): Whether the lookup should be blocking. Returns: Optional[Dict[str, str]]: The Regional Access Boundary information returned by the lookup API, or None if the lookup failed. @@ -456,7 +490,9 @@ def _lookup_regional_access_boundary( headers: Dict[str, str] = {} self._apply(headers) self._rab_manager.apply_headers(headers) - return _client._lookup_regional_access_boundary(request, url, headers=headers) + return _client._lookup_regional_access_boundary( + request, url, headers=headers, blocking=blocking + ) @abc.abstractmethod def _build_regional_access_boundary_lookup_url( diff --git a/packages/google-auth/google/oauth2/_client.py b/packages/google-auth/google/oauth2/_client.py index 4d76ed9a9a2b..751ed2420581 100644 --- a/packages/google-auth/google/oauth2/_client.py +++ b/packages/google-auth/google/oauth2/_client.py @@ -43,6 +43,7 @@ _JSON_CONTENT_TYPE = "application/json" _JWT_GRANT_TYPE = "urn:ietf:params:oauth:grant-type:jwt-bearer" _REFRESH_GRANT_TYPE = "refresh_token" +_BLOCKING_REGIONAL_ACCESS_BOUNDARY_LOOKUP_TIMEOUT = 3 def _handle_error_response(response_data, retryable_error): @@ -517,7 +518,7 @@ def refresh_grant( return _handle_refresh_grant_response(response_data, refresh_token) -def _lookup_regional_access_boundary(request, url, headers=None): +def _lookup_regional_access_boundary(request, url, headers=None, blocking=False): """Implements the global lookup of a credential Regional Access Boundary. For the lookup, we send a request to the global lookup endpoint and then parse the response. Service account credentials, workload identity @@ -527,6 +528,7 @@ def _lookup_regional_access_boundary(request, url, headers=None): HTTP requests. url (str): The Regional Access Boundary lookup url. headers (Optional[Mapping[str, str]]): The headers for the request. + blocking (bool): Whether the lookup should be blocking. Returns: Optional[Mapping[str,list|str]]: A dictionary containing "locations" as a list of allowed locations as strings and @@ -541,7 +543,7 @@ def _lookup_regional_access_boundary(request, url, headers=None): """ response_data = _lookup_regional_access_boundary_request( - request, url, headers=headers + request, url, headers=headers, blocking=blocking ) if response_data is None: # Error was already logged by _lookup_regional_access_boundary_request @@ -557,7 +559,7 @@ def _lookup_regional_access_boundary(request, url, headers=None): def _lookup_regional_access_boundary_request( - request, url, can_retry=True, headers=None + request, url, can_retry=True, headers=None, blocking=False ): """Makes a request to the Regional Access Boundary lookup endpoint. @@ -567,6 +569,7 @@ def _lookup_regional_access_boundary_request( url (str): The Regional Access Boundary lookup url. can_retry (bool): Enable or disable request retry behavior. Defaults to true. headers (Optional[Mapping[str, str]]): The headers for the request. + blocking (bool): Whether the lookup should be blocking. Returns: Optional[Mapping[str, str]]: The JSON-decoded response data on success, or None on failure. @@ -576,7 +579,7 @@ def _lookup_regional_access_boundary_request( response_data, retryable_error, ) = _lookup_regional_access_boundary_request_no_throw( - request, url, can_retry, headers + request, url, can_retry, headers, blocking ) if not response_status_ok: _LOGGER.warning( @@ -589,7 +592,7 @@ def _lookup_regional_access_boundary_request( def _lookup_regional_access_boundary_request_no_throw( - request, url, can_retry=True, headers=None + request, url, can_retry=True, headers=None, blocking=False ): """Makes a request to the Regional Access Boundary lookup endpoint. This function doesn't throw on response errors. @@ -600,6 +603,7 @@ def _lookup_regional_access_boundary_request_no_throw( url (str): The Regional Access Boundary lookup url. can_retry (bool): Enable or disable request retry behavior. Defaults to true. headers (Optional[Mapping[str, str]]): The headers for the request. + blocking (bool): Whether the lookup should be blocking. Returns: Tuple(bool, Mapping[str, str], Optional[bool]): A boolean indicating @@ -611,9 +615,12 @@ def _lookup_regional_access_boundary_request_no_throw( response_data = {} retryable_error = False - retries = _exponential_backoff.ExponentialBackoff(total_attempts=6) + timeout = _BLOCKING_REGIONAL_ACCESS_BOUNDARY_LOOKUP_TIMEOUT if blocking else None + total_attempts = 1 if blocking else 6 + retries = _exponential_backoff.ExponentialBackoff(total_attempts=total_attempts) + for _ in retries: - response = request(method="GET", url=url, headers=headers) + response = request(method="GET", url=url, headers=headers, timeout=timeout) response_body = ( response.data.decode("utf-8") if hasattr(response.data, "decode") diff --git a/packages/google-auth/google/oauth2/credentials.py b/packages/google-auth/google/oauth2/credentials.py index ae60223b455e..049c53da50f1 100644 --- a/packages/google-auth/google/oauth2/credentials.py +++ b/packages/google-auth/google/oauth2/credentials.py @@ -39,6 +39,7 @@ from google.auth import _cloud_sdk from google.auth import _helpers +from google.auth import _regional_access_boundary_utils from google.auth import credentials from google.auth import exceptions from google.auth import metrics @@ -54,7 +55,11 @@ _GOOGLE_OAUTH2_TOKEN_INFO_ENDPOINT = "https://oauth2.googleapis.com/tokeninfo" -class Credentials(credentials.ReadOnlyScoped, credentials.CredentialsWithQuotaProject): +class Credentials( + credentials.CredentialsWithRegionalAccessBoundary, + credentials.ReadOnlyScoped, + credentials.CredentialsWithQuotaProject, +): """Credentials using OAuth 2.0 access and refresh tokens. The credentials are considered immutable except the tokens and the token @@ -202,6 +207,10 @@ def __setstate__(self, d): self._refresh_worker = None self._use_non_blocking_refresh = d.get("_use_non_blocking_refresh", False) self._account = d.get("_account", "") + self._rab_manager = d.get("_rab_manager") or ( + _regional_access_boundary_utils._RegionalAccessBoundaryManager() + ) + self._use_blocking_regional_access_boundary_lookup = False @property def refresh_token(self): @@ -353,8 +362,10 @@ def with_universe_domain(self, universe_domain): def _metric_header_for_usage(self): return metrics.CRED_TYPE_USER - @_helpers.copy_docstring(credentials.Credentials) - def refresh(self, request): + def _build_regional_access_boundary_lookup_url(self, request=None): + return None + + def _perform_refresh_token(self, request): if self._universe_domain != credentials.DEFAULT_UNIVERSE_DOMAIN: raise exceptions.RefreshError( "User credential refresh is only supported in the default " diff --git a/packages/google-auth/tests/compute_engine/test_credentials.py b/packages/google-auth/tests/compute_engine/test_credentials.py index 4220bc84c2c7..c110a8a20e6d 100644 --- a/packages/google-auth/tests/compute_engine/test_credentials.py +++ b/packages/google-auth/tests/compute_engine/test_credentials.py @@ -451,6 +451,14 @@ def test_refresh_with_agent_identity_opt_out_or_not_agent( kwargs = mock_metadata_get.call_args[1] assert "bindCertificateFingerprint" not in kwargs.get("params", {}) + def test_with_blocking_regional_access_boundary_lookup(self): + creds = self.credentials + assert not creds._rab_manager._use_blocking_regional_access_boundary_lookup + + new_creds = creds.with_blocking_regional_access_boundary_lookup() + assert new_creds._rab_manager._use_blocking_regional_access_boundary_lookup + assert new_creds is not creds + class TestIDTokenCredentials(object): credentials = None diff --git a/packages/google-auth/tests/oauth2/test__client.py b/packages/google-auth/tests/oauth2/test__client.py index ad510f903d55..66a7ed19861d 100644 --- a/packages/google-auth/tests/oauth2/test__client.py +++ b/packages/google-auth/tests/oauth2/test__client.py @@ -654,7 +654,9 @@ def test_lookup_regional_access_boundary(): assert response["encodedLocations"] == "0x80080000000000" assert response["locations"] == ["us-central1", "us-east1"] - mock_request.assert_called_once_with(method="GET", url=url, headers=headers) + mock_request.assert_called_once_with( + method="GET", url=url, headers=headers, timeout=None + ) def test_lookup_regional_access_boundary_error(): @@ -672,7 +674,9 @@ def test_lookup_regional_access_boundary_error(): ) assert result is None - mock_request.assert_called_with(method="GET", url=url, headers=headers) + mock_request.assert_called_with( + method="GET", url=url, headers=headers, timeout=None + ) @pytest.mark.parametrize( @@ -697,7 +701,9 @@ def test_lookup_regional_access_boundary_non_retryable_error(status_code): ) assert result is None # Non-retryable errors should only be called once. - mock_request.assert_called_once_with(method="GET", url=url, headers=headers) + mock_request.assert_called_once_with( + method="GET", url=url, headers=headers, timeout=None + ) def test_lookup_regional_access_boundary_internal_failure_and_retry_failure_error(): @@ -777,5 +783,30 @@ def test_lookup_regional_access_boundary_with_headers(): ) mock_request.assert_called_once_with( - method="GET", url="http://example.com", headers=headers + method="GET", url="http://example.com", headers=headers, timeout=None + ) + + +def test_lookup_regional_access_boundary_blocking(): + response_data = { + "locations": ["us-central1"], + "encodedLocations": "0xABC", + } + + mock_response = mock.create_autospec(transport.Response, instance=True) + mock_response.status = http_client.OK + mock_response.data = json.dumps(response_data).encode("utf-8") + + mock_request = mock.create_autospec(transport.Request) + mock_request.return_value = mock_response + + url = "http://example.com" + headers = {"Authorization": "Bearer access_token"} + response = _client._lookup_regional_access_boundary( + mock_request, url, headers=headers, blocking=True + ) + + assert response["encodedLocations"] == "0xABC" + mock_request.assert_called_once_with( + method="GET", url=url, headers=headers, timeout=3 ) diff --git a/packages/google-auth/tests/oauth2/test_credentials.py b/packages/google-auth/tests/oauth2/test_credentials.py index f22a0a2ad918..7250ce9f0cf5 100644 --- a/packages/google-auth/tests/oauth2/test_credentials.py +++ b/packages/google-auth/tests/oauth2/test_credentials.py @@ -96,6 +96,14 @@ def test_default_state(self): assert credentials.rapt_token == self.RAPT_TOKEN assert credentials.refresh_handler is None + def test_with_blocking_regional_access_boundary_lookup(self): + creds = self.make_credentials() + assert not creds._rab_manager._use_blocking_regional_access_boundary_lookup + + new_creds = creds.with_blocking_regional_access_boundary_lookup() + assert new_creds._rab_manager._use_blocking_regional_access_boundary_lookup + assert new_creds is not creds + def test_get_cred_info(self): credentials = self.make_credentials() credentials._account = "fake-account" diff --git a/packages/google-auth/tests/test__regional_access_boundary_utils.py b/packages/google-auth/tests/test__regional_access_boundary_utils.py index f3002d204a6a..0debc89837a5 100644 --- a/packages/google-auth/tests/test__regional_access_boundary_utils.py +++ b/packages/google-auth/tests/test__regional_access_boundary_utils.py @@ -223,6 +223,24 @@ def test_apply_headers_removes_header_if_empty(self): creds._rab_manager.apply_headers(headers) assert headers == {} + def test_with_blocking_regional_access_boundary_lookup(self): + creds = CredentialsImpl() + assert not creds._rab_manager._use_blocking_regional_access_boundary_lookup + + new_creds = creds.with_blocking_regional_access_boundary_lookup() + assert new_creds._rab_manager._use_blocking_regional_access_boundary_lookup + + def test_with_regional_access_boundary(self): + creds = CredentialsImpl() + seed = { + "encodedLocations": "0xABC", + "expiry": _helpers.utcnow() + datetime.timedelta(hours=1), + } + new_creds = creds._with_regional_access_boundary(seed) + assert new_creds._rab_manager._data.encoded_locations == "0xABC" + assert new_creds._rab_manager._data.expiry == seed["expiry"] + assert new_creds._rab_manager._data.cooldown_expiry is None + def test_copy_regional_access_boundary_state(self): source_creds = CredentialsImpl() snapshot = _regional_access_boundary_utils._RegionalAccessBoundaryData( @@ -293,59 +311,99 @@ def test_maybe_start_refresh_handles_url_parse_errors( ) mock_start_refresh.assert_called_once_with(creds, request, creds._rab_manager) - @mock.patch("google.oauth2._client._lookup_regional_access_boundary") - @mock.patch.object(CredentialsImpl, "_build_regional_access_boundary_lookup_url") - def test_lookup_regional_access_boundary_success( - self, mock_build_url, mock_lookup_rab - ): + @mock.patch( + "google.auth._regional_access_boundary_utils._RegionalAccessBoundaryManager.start_blocking_refresh" + ) + def test_maybe_start_refresh_blocking(self, mock_start_blocking_refresh): + creds = CredentialsImpl() + creds._rab_manager._use_blocking_regional_access_boundary_lookup = True + request = mock.Mock() + with mock.patch.dict( + os.environ, + {environment_vars.GOOGLE_AUTH_TRUST_BOUNDARY_ENABLED: "true"}, + ): + creds._maybe_start_regional_access_boundary_refresh( + request, "http://example.com" + ) + mock_start_blocking_refresh.assert_called_once_with(creds, request) + + def test_start_blocking_refresh_success(self): creds = CredentialsImpl() - creds.token = "token" request = mock.Mock() - mock_build_url.return_value = "http://rab.example.com" - mock_lookup_rab.return_value = {"encodedLocations": "success"} - result = creds._lookup_regional_access_boundary(request) + with mock.patch.object( + creds, + "_lookup_regional_access_boundary", + return_value={"encodedLocations": "0xABC"}, + ) as mock_lookup: + creds._rab_manager.start_blocking_refresh(creds, request) + + mock_lookup.assert_called_once_with(request, True) + assert creds._rab_manager._data.encoded_locations == "0xABC" + + def test_start_blocking_refresh_failure(self): + creds = CredentialsImpl() + request = mock.Mock() + + with mock.patch.object( + creds, "_lookup_regional_access_boundary", side_effect=Exception("error") + ) as mock_lookup: + creds._rab_manager.start_blocking_refresh(creds, request) + + mock_lookup.assert_called_once_with(request, True) + assert creds._rab_manager._data.encoded_locations is None + assert creds._rab_manager._data.cooldown_expiry is not None + + @mock.patch("copy.deepcopy") + def test_start_refresh_deepcopy_failure(self, mock_deepcopy): + mock_deepcopy.side_effect = Exception("deepcopy error") + creds = CredentialsImpl() + request = mock.Mock() - mock_build_url.assert_called_once() - mock_lookup_rab.assert_called_once_with( - request, "http://rab.example.com", headers={"authorization": "Bearer token"} + creds._rab_manager.refresh_manager.start_refresh( + creds, request, creds._rab_manager ) - assert result == {"encodedLocations": "success"} - @mock.patch("google.oauth2._client._lookup_regional_access_boundary") - @mock.patch.object(CredentialsImpl, "_build_regional_access_boundary_lookup_url") - def test_lookup_regional_access_boundary_failure( - self, mock_build_url, mock_lookup_rab - ): + assert creds._rab_manager.refresh_manager._worker is None + + @mock.patch.object(CredentialsImpl, "_lookup_regional_access_boundary") + def test_lookup_regional_access_boundary_success(self, mock_lookup_rab): creds = CredentialsImpl() - creds.token = "token" request = mock.Mock() - mock_build_url.return_value = "http://rab.example.com" - mock_lookup_rab.return_value = None + rab_manager = _regional_access_boundary_utils._RegionalAccessBoundaryManager() - result = creds._lookup_regional_access_boundary(request) + mock_lookup_rab.return_value = { + "locations": ["us-east1"], + "encodedLocations": "0xABC123", + } - mock_build_url.assert_called_once() - mock_lookup_rab.assert_called_once_with( - request, "http://rab.example.com", headers={"authorization": "Bearer token"} + worker = _regional_access_boundary_utils._RegionalAccessBoundaryRefreshThread( + creds, request, rab_manager ) - assert result is None + worker.run() - @mock.patch("google.oauth2._client._lookup_regional_access_boundary") - @mock.patch.object(CredentialsImpl, "_build_regional_access_boundary_lookup_url") - def test_lookup_regional_access_boundary_null_url( - self, mock_build_url, mock_lookup_rab - ): + mock_lookup_rab.assert_called_once_with(request) + assert rab_manager._data.encoded_locations == "0xABC123" + assert rab_manager._data.expiry is not None + assert rab_manager._data.cooldown_expiry is None + + @mock.patch.object(CredentialsImpl, "_lookup_regional_access_boundary") + def test_lookup_regional_access_boundary_failure(self, mock_lookup_rab): creds = CredentialsImpl() - creds.token = "token" request = mock.Mock() - mock_build_url.return_value = None + rab_manager = _regional_access_boundary_utils._RegionalAccessBoundaryManager() - result = creds._lookup_regional_access_boundary(request) + mock_lookup_rab.return_value = None - mock_build_url.assert_called_once() - mock_lookup_rab.assert_not_called() - assert result is None + worker = _regional_access_boundary_utils._RegionalAccessBoundaryRefreshThread( + creds, request, rab_manager + ) + worker.run() + + mock_lookup_rab.assert_called_once_with(request) + assert rab_manager._data.encoded_locations is None + assert rab_manager._data.expiry is None + assert rab_manager._data.cooldown_expiry is not None def test_credentials_with_regional_access_boundary_initialization(self): creds = CredentialsImpl()