diff --git a/google/auth/compute_engine/_metadata.py b/google/auth/compute_engine/_metadata.py index 2861192d0..cba426fb2 100644 --- a/google/auth/compute_engine/_metadata.py +++ b/google/auth/compute_engine/_metadata.py @@ -32,9 +32,16 @@ _LOGGER = logging.getLogger(__name__) -_METADATA_ROOT = "http://{}/computeMetadata/v1/".format( - os.getenv(environment_vars.GCE_METADATA_ROOT, "metadata.google.internal") -) +# Environment variable GCE_METADATA_HOST is originally named +# GCE_METADATA_ROOT. For compatiblity reasons, here it checks +# the new variable first; if not set, the system falls back +# to the old variable. +_GCE_METADATA_HOST = os.getenv(environment_vars.GCE_METADATA_HOST, None) +if not _GCE_METADATA_HOST: + _GCE_METADATA_HOST = os.getenv( + environment_vars.GCE_METADATA_ROOT, "metadata.google.internal" + ) +_METADATA_ROOT = "http://{}/computeMetadata/v1/".format(_GCE_METADATA_HOST) # This is used to ping the metadata server, it avoids the cost of a DNS # lookup. diff --git a/google/auth/environment_vars.py b/google/auth/environment_vars.py index 6a596f22d..9c1367f52 100644 --- a/google/auth/environment_vars.py +++ b/google/auth/environment_vars.py @@ -40,9 +40,15 @@ # These two variables allow for customization of the addresses used when # contacting the GCE metadata service. +GCE_METADATA_HOST = "GCE_METADATA_HOST" GCE_METADATA_ROOT = "GCE_METADATA_ROOT" """Environment variable providing an alternate hostname or host:port to be -used for GCE metadata requests.""" +used for GCE metadata requests. + +This environment variable is originally named GCE_METADATA_ROOT. System will +check the new variable first; should there be no value present, +the system falls back to the old variable. +""" GCE_METADATA_IP = "GCE_METADATA_IP" """Environment variable providing an alternate ip:port to be used for ip-only diff --git a/tests/compute_engine/test__metadata.py b/tests/compute_engine/test__metadata.py index 8b5eece7a..d9b039a32 100644 --- a/tests/compute_engine/test__metadata.py +++ b/tests/compute_engine/test__metadata.py @@ -155,7 +155,27 @@ def test_get_success_text(): assert result == data -def test_get_success_custom_root(): +def test_get_success_custom_root_new_variable(): + request = make_request("{}", headers={"content-type": "application/json"}) + + fake_root = "another.metadata.service" + os.environ[environment_vars.GCE_METADATA_HOST] = fake_root + reload_module(_metadata) + + try: + _metadata.get(request, PATH) + finally: + del os.environ[environment_vars.GCE_METADATA_HOST] + reload_module(_metadata) + + request.assert_called_once_with( + method="GET", + url="http://{}/computeMetadata/v1/{}".format(fake_root, PATH), + headers=_metadata._METADATA_HEADERS, + ) + + +def test_get_success_custom_root_old_variable(): request = make_request("{}", headers={"content-type": "application/json"}) fake_root = "another.metadata.service"