Repository: GoogleCloudPlatform/google-auth-library-python
Pull request #148 (2017-03-24 21:33:33)

Allow customizing the GCE metadata service address via an env var.
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.)

The only bit of awkwardness here is really the test: since this is a top-level
statement, reloading is the only way to ensure it works.

PTAL @jonparrott

(and why yes I DID just completely copy-paste the commit description from the oauth2client PR.)

dhermes: @craigcitro Did you want to merge in both libraries? (It's pretty unlikely `oauth2client` will ever get another release.)
craigcitro: @dhermes Yep, I purposely sent the PR for both.

@jonparrott mentioned the same about oauth2client, then remembered that he had to add something silly (update the classifiers?), so it was an opportune time for me to sneak it in. 😉 
craigcitro: OK, updated to use the constants in `environment_vars` -- PTAL
jonparrott: @craigcitro LGTM, but Travis is mad about you not using 4 spaces for indentation. :)
craigcitro: As well he should be. Fixed.

(Do you guys autosquash or should I squash?)
jonparrott: > (Do you guys autosquash or should I squash?)

We autosquash

![crop-focus-successful-squash](https://cloud.githubusercontent.com/assets/250995/24315789/8e338dac-10a6-11e7-91d9-efc4e0b73730.jpg)




Patch:
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 import http_client
 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 @@ credentials."""
 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."""


Test patch:
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 @@ 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)