diff --git a/google/auth/compute_engine/_metadata.py b/google/auth/compute_engine/_metadata.py index ef726fe7a..c1fecd3b7 100644 --- a/google/auth/compute_engine/_metadata.py +++ b/google/auth/compute_engine/_metadata.py @@ -97,10 +97,19 @@ def _get_metadata_ip_root(use_mtls: bool): # Timeout in seconds to wait for the GCE metadata server when detecting the # GCE environment. try: - _METADATA_DEFAULT_TIMEOUT = int(os.getenv("GCE_METADATA_TIMEOUT", 3)) + _METADATA_DEFAULT_TIMEOUT = int(os.getenv(environment_vars.GCE_METADATA_TIMEOUT, 3)) except ValueError: # pragma: NO COVER _METADATA_DEFAULT_TIMEOUT = 3 +# The number of tries to perform when waiting for the GCE metadata server +# when detecting the GCE environment. +try: + _METADATA_DETECT_RETRIES = int( + os.getenv(environment_vars.GCE_METADATA_DETECT_RETRIES, 3) + ) +except ValueError: # pragma: NO COVER + _METADATA_DETECT_RETRIES = 3 + # This is used to disable checking for the GCE metadata server and directly # assuming it's not available. _NO_GCE_CHECK = os.getenv(environment_vars.NO_GCE_CHECK) == "true" @@ -177,7 +186,9 @@ def _prepare_request_for_mds(request, use_mtls=False) -> None: request.session.mount(f"https://{host}/", adapter) -def ping(request, timeout=_METADATA_DEFAULT_TIMEOUT, retry_count=3): +def ping( + request, timeout=_METADATA_DEFAULT_TIMEOUT, retry_count=_METADATA_DETECT_RETRIES +): """Checks to see if the metadata server is available. Args: diff --git a/google/auth/environment_vars.py b/google/auth/environment_vars.py index 587bdcfdb..ca041ca16 100644 --- a/google/auth/environment_vars.py +++ b/google/auth/environment_vars.py @@ -60,6 +60,16 @@ """Environment variable providing an alternate ip:port to be used for ip-only GCE metadata requests.""" +GCE_METADATA_TIMEOUT = "GCE_METADATA_TIMEOUT" +"""Environment variable defining the timeout in seconds to wait for the +GCE metadata server when detecting the GCE environment. +""" + +GCE_METADATA_DETECT_RETRIES = "GCE_METADATA_DETECT_RETRIES" +"""Environment variable representing the number of retries that should be +attempted on metadata lookup. +""" + NO_GCE_CHECK = "NO_GCE_CHECK" """Environment variable controlling whether to check if running on GCE or not. diff --git a/tests/compute_engine/test__metadata.py b/tests/compute_engine/test__metadata.py index fd4f8255b..7bd7bc6ed 100644 --- a/tests/compute_engine/test__metadata.py +++ b/tests/compute_engine/test__metadata.py @@ -201,6 +201,23 @@ def test_ping_success_custom_root(mock_metrics_header_value): ) +@mock.patch("google.auth.metrics.mds_ping", return_value=MDS_PING_METRICS_HEADER_VALUE) +def test_ping_failure_custom_retry(mock_metrics_header_value): + request = make_request("") + request.side_effect = exceptions.TransportError() + + os.environ[environment_vars.GCE_METADATA_DETECT_RETRIES] = "10" + importlib.reload(_metadata) + + try: + _metadata.ping(request) + finally: + del os.environ[environment_vars.GCE_METADATA_DETECT_RETRIES] + importlib.reload(_metadata) + + assert request.call_count == 10 + + def test_get_success_json(): key, value = "foo", "bar"