Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
15 changes: 13 additions & 2 deletions google/auth/compute_engine/_metadata.py
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Expand Down Expand Up @@ -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:
Expand Down
10 changes: 10 additions & 0 deletions google/auth/environment_vars.py
Original file line number Diff line number Diff line change
Expand Up @@ -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.

Expand Down
17 changes: 17 additions & 0 deletions tests/compute_engine/test__metadata.py
Original file line number Diff line number Diff line change
Expand Up @@ -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"

Expand Down