Allow customizing the GCE metadata service address via an env var. (#148)
The goal here is to make it possible for a user of a binary that depends on
this library (eg the google cloud SDK) to be able to customize where it looks
for the GCE metadata service. (An adventurous user can already customize the
GCE metadata service location via the existing global vars in this library.)
diff --git a/google/auth/compute_engine/_metadata.py b/google/auth/compute_engine/_metadata.py
index e5abf20..b35775a 100644
--- a/google/auth/compute_engine/_metadata.py
+++ b/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}
diff --git a/google/auth/environment_vars.py b/google/auth/environment_vars.py
index b4ed2b2..0110e6a 100644
--- a/google/auth/environment_vars.py
+++ b/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."""
diff --git a/tests/compute_engine/test__metadata.py b/tests/compute_engine/test__metadata.py
index 6a45c42..0adf882 100644
--- a/tests/compute_engine/test__metadata.py
+++ b/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 @@
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 @@
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)