Allow customizing the GCE metadata service address via an env var.

google/auth/compute_engine/_metadata.py

@@ -26,15 +26,18 @@
 from six.moves.urllib import parse as urlparse
 
 from google.auth import _helpers
+from google.auth import environment_vars
 from google.auth import exceptions
 
 _LOGGER = logging.getLogger(__name__)
 
-_METADATA_ROOT = 'http://metadata.google.internal/computeMetadata/v1/'
+_METADATA_ROOT = 'http://{}/computeMetadata/v1/'.format(
+    os.getenv(environment_vars.GCE_METADATA_ROOT, 'metadata.google.internal'))
 
 # This is used to ping the metadata server, it avoids the cost of a DNS
 # lookup.
-_METADATA_IP_ROOT = 'http://169.254.169.254'
+_METADATA_IP_ROOT = 'http://{}'.format(
+    os.getenv(environment_vars.GCE_METADATA_IP, '169.254.169.254'))
 _METADATA_FLAVOR_HEADER = 'metadata-flavor'
 _METADATA_FLAVOR_VALUE = 'Google'
 _METADATA_HEADERS = {_METADATA_FLAVOR_HEADER: _METADATA_FLAVOR_VALUE}

google/auth/environment_vars.py

@@ -37,3 +37,13 @@
 CLOUD_SDK_CONFIG_DIR = 'CLOUDSDK_CONFIG'
 """Environment variable defines the location of Google Cloud SDK's config
 files."""
+
+# These two variables allow for customization of the addresses used when
+# contacting the GCE metadata service.
+GCE_METADATA_ROOT = 'GCE_METADATA_ROOT'
+"""Environment variable providing an alternate hostname or host:port to be
+used for GCE metadata requests."""
+
+GCE_METADATA_IP = 'GCE_METADATA_IP'
+"""Environment variable providing an alternate ip:port to be used for ip-only
+GCE metadata requests."""

tests/compute_engine/test__metadata.py

@@ -14,12 +14,15 @@
 
 import datetime
 import json
+import os
 
 import mock
 import pytest
 from six.moves import http_client
+from six.moves import reload_module
 
 from google.auth import _helpers
+from google.auth import environment_vars
 from google.auth import exceptions
 from google.auth.compute_engine import _metadata
 
@@ -67,6 +70,26 @@ def test_ping_failure_connection_failed(mock_request):
     assert not _metadata.ping(request_mock)
 
 
+def test_ping_success_custom_root(mock_request):
+    request_mock = mock_request('', headers=_metadata._METADATA_HEADERS)
+
+    fake_ip = '1.2.3.4'
+    os.environ[environment_vars.GCE_METADATA_IP] = fake_ip
+    reload_module(_metadata)
+
+    try:
+        assert _metadata.ping(request_mock)
+    finally:
+        del os.environ[environment_vars.GCE_METADATA_IP]
+        reload_module(_metadata)
+
+    request_mock.assert_called_once_with(
+        method='GET',
+        url='http://' + fake_ip,
+        headers=_metadata._METADATA_HEADERS,
+        timeout=_metadata._METADATA_DEFAULT_TIMEOUT)
+
+
 def test_get_success_json(mock_request):
     key, value = 'foo', 'bar'
 
@@ -96,6 +119,26 @@ def test_get_success_text(mock_request):
     assert result == data
 
 
+def test_get_success_custom_root(mock_request):
+    request_mock = mock_request(
+        '{}', headers={'content-type': 'application/json'})
+
+    fake_root = 'another.metadata.service'
+    os.environ[environment_vars.GCE_METADATA_ROOT] = fake_root
+    reload_module(_metadata)
+
+    try:
+        _metadata.get(request_mock, PATH)
+    finally:
+        del os.environ[environment_vars.GCE_METADATA_ROOT]
+        reload_module(_metadata)
+
+    request_mock.assert_called_once_with(
+        method='GET',
+        url='http://{}/computeMetadata/v1/{}'.format(fake_root, PATH),
+        headers=_metadata._METADATA_HEADERS)
+
+
 def test_get_failure(mock_request):
     request_mock = mock_request(
         'Metadata error', status=http_client.NOT_FOUND)