Merge branch 'main' into add_cryptography_dependency

Author: daniel-sanche

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:

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.
 

samples/cloud-client/snippets/requirements.txt

@@ -1,6 +1,6 @@
-google-cloud-compute==1.40.0
+google-cloud-compute==1.41.0
 google-cloud-storage==3.7.0
-google-auth==2.45.0
+google-auth==2.47.0
 pytest===8.4.2; python_version == '3.9'
 pytest==9.0.2; python_version > '3.9'
 boto3>=1.26.0

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"