feat: make number of metadata detection retries configurable

Author: judahrand

google/auth/compute_engine/_metadata.py

@@ -58,6 +58,15 @@
 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:

google/auth/environment_vars.py

@@ -60,6 +60,11 @@
 """Environment variable providing an alternate ip:port to be used for ip-only
 GCE metadata requests."""
 
+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.
 

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"