diff --git a/google/auth/compute_engine/_metadata.py b/google/auth/compute_engine/_metadata.py index 108cbfe93..9a54c9158 100644 --- a/google/auth/compute_engine/_metadata.py +++ b/google/auth/compute_engine/_metadata.py @@ -54,10 +54,19 @@ # 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 + # Detect GCE Residency _GOOGLE = "Google" _GCE_PRODUCT_NAME_FILE = "/sys/class/dmi/id/product_name" @@ -100,7 +109,9 @@ def detect_gce_residency_linux(): return content.startswith(_GOOGLE) -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 81f31571e..1cc8c5ee8 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. +""" + GOOGLE_API_USE_CLIENT_CERTIFICATE = "GOOGLE_API_USE_CLIENT_CERTIFICATE" """Environment variable controlling whether to use client certificate or not. diff --git a/tests/compute_engine/test__metadata.py b/tests/compute_engine/test__metadata.py index 60ae355ac..2cd077c21 100644 --- a/tests/compute_engine/test__metadata.py +++ b/tests/compute_engine/test__metadata.py @@ -160,6 +160,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"