Refactor to avoid monkeypatching requests

os_collect_config/collect.py

@@ -19,6 +19,7 @@ import subprocess
 
 from openstack.common import log
 from os_collect_config import cache
+from os_collect_config import common
 from os_collect_config import ec2
 from oslo.config import cfg
 
@@ -45,11 +46,11 @@ def setup_conf():
     CONF.register_cli_opts(opts)
 
 
-def __main__():
+def __main__(ec2_requests=common.requests):
     setup_conf()
     CONF(prog="os-collect-config")
     log.setup("os-collect-config")
-    ec2_content = ec2.collect()
+    ec2_content = ec2.CollectEc2(requests_impl=ec2_requests).collect()
 
     if CONF.command:
         (changed, ec2_path) = cache.store('ec2', ec2_content)

os_collect_config/common.py

@@ -0,0 +1,3 @@
+import requests
+
+__all__ = ['requests']

os_collect_config/ec2.py

@@ -14,9 +14,9 @@
 # limitations under the License.
 
 from oslo.config import cfg
-import requests
 
 from openstack.common import log
+from os_collect_config import common
 from os_collect_config import exc
 
 EC2_METADATA_URL = 'http://169.254.169.254/latest/meta-data'
@@ -29,30 +29,31 @@ opts = [
 ]
 
 
-def _fetch_metadata(fetch_url, session):
+class CollectEc2(object):
-    try:
+    def __init__(self, requests_impl=common.requests):
-        r = session.get(fetch_url)
+        self._requests_impl = requests_impl
-        r.raise_for_status()
+        self.session = requests_impl.Session()
-    except (requests.HTTPError,
-            requests.ConnectionError,
-            requests.Timeout) as e:
-        log.getLogger(__name__).warn(e)
-        raise exc.Ec2MetadataNotAvailable
-    content = r.text
-    if fetch_url[-1] == '/':
-        new_content = {}
-        for subkey in content.split("\n"):
-            if '=' in subkey:
-                subkey = subkey[:subkey.index('=')] + '/'
-            sub_fetch_url = fetch_url + subkey
-            if subkey[-1] == '/':
-                subkey = subkey[:-1]
-            new_content[subkey] = _fetch_metadata(sub_fetch_url, session)
-        content = new_content
-    return content
 
+    def _fetch_metadata(self, fetch_url):
+        try:
+            r = self.session.get(fetch_url)
+            r.raise_for_status()
+        except self._requests_impl.exceptions.RequestException as e:
+            log.getLogger(__name__).warn(e)
+            raise exc.Ec2MetadataNotAvailable
+        content = r.text
+        if fetch_url[-1] == '/':
+            new_content = {}
+            for subkey in content.split("\n"):
+                if '=' in subkey:
+                    subkey = subkey[:subkey.index('=')] + '/'
+                sub_fetch_url = fetch_url + subkey
+                if subkey[-1] == '/':
+                    subkey = subkey[:-1]
+                new_content[subkey] = self._fetch_metadata(sub_fetch_url)
+            content = new_content
+        return content
 
-def collect():
+    def collect(self):
-    root_url = '%s/' % (CONF.ec2.metadata_url)
+        root_url = '%s/' % (CONF.ec2.metadata_url)
-    session = requests.Session()
+        return self._fetch_metadata(root_url)
-    return _fetch_metadata(root_url, session)

os_collect_config/tests/test_collect.py

@@ -25,12 +25,6 @@ from os_collect_config.tests import test_ec2
 
 
 class TestCollect(testtools.TestCase):
-    def setUp(self):
-        super(TestCollect, self).setUp()
-        self.useFixture(
-            fixtures.MonkeyPatch(
-                'requests.Session', test_ec2.FakeSession))
-
     def tearDown(self):
         super(TestCollect, self).tearDown()
         cfg.CONF.reset()
@@ -67,7 +61,7 @@ class TestCollect(testtools.TestCase):
 
         self.useFixture(fixtures.MonkeyPatch('subprocess.call', fake_call))
 
-        collect.__main__()
+        collect.__main__(ec2_requests=test_ec2.FakeRequests)
 
         self.assertTrue(self.called_fake_call)
 
@@ -82,7 +76,7 @@ class TestCollect(testtools.TestCase):
         output = self.useFixture(fixtures.ByteStream('stdout'))
         self.useFixture(
             fixtures.MonkeyPatch('sys.stdout', output.stream))
-        collect.__main__()
+        collect.__main__(ec2_requests=test_ec2.FakeRequests)
         out_struct = json.loads(output.stream.getvalue())
         self.assertThat(out_struct, matchers.IsInstance(dict))
         self.assertIn('ec2', out_struct)

os_collect_config/tests/test_ec2.py

@@ -58,36 +58,40 @@ class FakeResponse(dict):
         pass
 
 
-class FakeSession(object):
+class FakeRequests(object):
-    def get(self, url):
+    exceptions = requests.exceptions
-        url = urlparse.urlparse(url)
 
-        if url.path == '/latest/meta-data/':
+    class Session(object):
-            # Remove keys which have anything after /
+        def get(self, url):
-            ks = [x for x in META_DATA.keys() if ('/' not in x
+            url = urlparse.urlparse(url)
-                                                  or not len(x.split('/')[1]))]
-            return FakeResponse("\n".join(ks))
 
-        path = url.path
+            if url.path == '/latest/meta-data/':
-        path = path.replace('/latest/meta-data/', '')
+                # Remove keys which have anything after /
-        return FakeResponse(META_DATA[path])
+                ks = [x for x in META_DATA.keys() if (
+                    '/' not in x or not len(x.split('/')[1]))]
+                return FakeResponse("\n".join(ks))
+
+            path = url.path
+            path = path.replace('/latest/meta-data/', '')
+            return FakeResponse(META_DATA[path])
 
 
-class FakeFailSession(object):
+class FakeFailRequests(object):
-    def get(self, url):
+    exceptions = requests.exceptions
-        raise requests.exceptions.HTTPError(403, 'Forbidden')
+
+    class Session(object):
+        def get(self, url):
+            raise requests.exceptions.HTTPError(403, 'Forbidden')
 
 
-class TestCollect(testtools.TestCase):
+class TestEc2(testtools.TestCase):
     def setUp(self):
-        super(TestCollect, self).setUp()
+        super(TestEc2, self).setUp()
         self.log = self.useFixture(fixtures.FakeLogger())
 
     def test_collect_ec2(self):
-        self.useFixture(
-            fixtures.MonkeyPatch('requests.Session', FakeSession))
         collect.setup_conf()
-        ec2_md = ec2.collect()
+        ec2_md = ec2.CollectEc2(requests_impl=FakeRequests).collect()
         self.assertThat(ec2_md, matchers.IsInstance(dict))
 
         for k in ('public-ipv4', 'instance-id', 'hostname'):
@@ -103,9 +107,7 @@ class TestCollect(testtools.TestCase):
         self.assertEquals('', self.log.output)
 
     def test_collect_ec2_fail(self):
-        self.useFixture(
-            fixtures.MonkeyPatch(
-                'requests.Session', FakeFailSession))
         collect.setup_conf()
-        self.assertRaises(exc.Ec2MetadataNotAvailable, ec2.collect)
+        collect_ec2 = ec2.CollectEc2(requests_impl=FakeFailRequests)
+        self.assertRaises(exc.Ec2MetadataNotAvailable, collect_ec2.collect)
         self.assertIn('Forbidden', self.log.output)