From 9722b3206f8e6e2c326f039513bac176487fabe3 Mon Sep 17 00:00:00 2001
From: Craig Citro <craigcitro@gmail.com>
Date: Fri, 19 Dec 2014 11:34:07 -0800
Subject: [PATCH] Add a timeout for GCE detection.

This is a temporary workaround for #93, though may end up being the final fix.

I did some test cleanup while I was here, switching to `mock`.
---
 oauth2client/client.py     | 17 ++++++++++++---
 tests/test_oauth2client.py | 44 +++++++++++---------------------------
 tox.ini                    |  3 ++-
 3 files changed, 28 insertions(+), 36 deletions(-)

diff --git a/oauth2client/client.py b/oauth2client/client.py
index b7cefc549..f45edc59f 100644
--- a/oauth2client/client.py
+++ b/oauth2client/client.py
@@ -26,6 +26,7 @@
 import json
 import logging
 import os
+import socket
 import sys
 import time
 import six
@@ -930,12 +931,22 @@ def _detect_gce_environment(urlopen=None):
           Compute Engine.
   """
   urlopen = urlopen or urllib.request.urlopen
-
+  # Note: the explicit `timeout` below is a workaround. The underlying
+  # issue is that resolving an unknown host on some networks will take
+  # 20-30 seconds; making this timeout short fixes the issue, but
+  # could lead to false negatives in the event that we are on GCE, but
+  # the metadata resolution was particularly slow. The latter case is
+  # "unlikely".
   try:
-    response = urlopen('http://metadata.google.internal')
+    response = urlopen('http://metadata.google.internal/', timeout=1)
     return any('Metadata-Flavor: Google' in header
                for header in response.info().headers)
-  except urllib.error.URLError:
+  except socket.timeout:
+    logger.info('Timeout attempting to reach GCE metadata service.')
+    return False
+  except urllib.error.URLError as e:
+    if isinstance(getattr(e, 'reason', None), socket.timeout):
+      logger.info('Timeout attempting to reach GCE metadata service.')
     return False
 
 
diff --git a/tests/test_oauth2client.py b/tests/test_oauth2client.py
index 82ddb92cd..6b719c846 100644
--- a/tests/test_oauth2client.py
+++ b/tests/test_oauth2client.py
@@ -25,13 +25,11 @@
 import base64
 import datetime
 import json
-try:
-  from mox3 import mox
-except ImportError:
-  import mox
 import os
 import time
 import unittest
+
+import mock
 import six
 from six.moves import urllib
 
@@ -208,37 +206,19 @@ def test_get_environment_gae_local(self):
 
   def test_get_environment_gce_production(self):
     os.environ['SERVER_SOFTWARE'] = ''
-    mockResponse = MockResponse(['Metadata-Flavor: Google\r\n'])
-
-    m = mox.Mox()
-
-    urllib2_urlopen = m.CreateMock(object)
-    urllib2_urlopen.__call__(('http://metadata.google.internal'
-                             )).AndReturn(mockResponse)
-
-    m.ReplayAll()
-
-    self.assertEqual('GCE_PRODUCTION', _get_environment(urllib2_urlopen))
-
-    m.UnsetStubs()
-    m.VerifyAll()
+    with mock.patch.object(urllib.request, 'urlopen') as urlopen:
+      urlopen.return_value = MockResponse(['Metadata-Flavor: Google\r\n'])
+      self.assertEqual('GCE_PRODUCTION', _get_environment())
+      urlopen.assert_called_once_with(
+          'http://metadata.google.internal/', timeout=1)
 
   def test_get_environment_unknown(self):
     os.environ['SERVER_SOFTWARE'] = ''
-    mockResponse = MockResponse([])
-
-    m = mox.Mox()
-
-    urllib2_urlopen = m.CreateMock(object)
-    urllib2_urlopen.__call__(('http://metadata.google.internal'
-                             )).AndReturn(mockResponse)
-
-    m.ReplayAll()
-
-    self.assertEqual(DEFAULT_ENV_NAME, _get_environment(urllib2_urlopen))
-
-    m.UnsetStubs()
-    m.VerifyAll()
+    with mock.patch.object(urllib.request, 'urlopen') as urlopen:
+      urlopen.return_value = MockResponse([])
+      self.assertEqual(DEFAULT_ENV_NAME, _get_environment())
+      urlopen.assert_called_once_with(
+          'http://metadata.google.internal/', timeout=1)
 
   def test_get_environment_variable_file(self):
     environment_variable_file = datafile(
diff --git a/tox.ini b/tox.ini
index b122f079a..d2ff4c534 100644
--- a/tox.ini
+++ b/tox.ini
@@ -4,6 +4,7 @@ envlist = py26,py27,py33,py34,pypy,cover
 [testenv]
 basedeps = keyring
            mox3
+           mock
            pycrypto==2.6
            django>=1.5,<1.6
            webtest
@@ -11,7 +12,7 @@ basedeps = keyring
 deps = {[testenv]basedeps}
        pyopenssl==0.14
 setenv = PYTHONPATH=../google_appengine
-commands = nosetests --ignore-files=test_appengine\.py
+commands = nosetests --ignore-files=test_appengine\.py {posargs}
 
 # whitelist
 branches: