diff options
| -rw-r--r-- | tests/unittests/test_datasource/test_gce.py | 47 |
1 files changed, 12 insertions, 35 deletions
diff --git a/tests/unittests/test_datasource/test_gce.py b/tests/unittests/test_datasource/test_gce.py index 1fb100f7..98b68f09 100644 --- a/tests/unittests/test_datasource/test_gce.py +++ b/tests/unittests/test_datasource/test_gce.py @@ -55,8 +55,8 @@ MD_URL_RE = re.compile( r'http://metadata.google.internal./computeMetadata/v1/.*') -def _new_request_callback(gce_meta=None): - if not gce_meta: +def _set_mock_metadata(gce_meta=None): + if gce_meta is None: gce_meta = GCE_META def _request_callback(method, uri, headers): @@ -70,9 +70,10 @@ def _new_request_callback(gce_meta=None): else: return (404, headers, '') - return _request_callback + httpretty.register_uri(httpretty.GET, MD_URL_RE, body=_request_callback) +@httpretty.activate class TestDataSourceGCE(test_helpers.HttprettyTestCase): def setUp(self): @@ -81,23 +82,16 @@ class TestDataSourceGCE(test_helpers.HttprettyTestCase): helpers.Paths({})) super(TestDataSourceGCE, self).setUp() - @httpretty.activate def test_connection(self): - httpretty.register_uri( - httpretty.GET, MD_URL_RE, - body=_new_request_callback()) - + _set_mock_metadata() success = self.ds.get_data() self.assertTrue(success) req_header = httpretty.last_request().headers self.assertDictContainsSubset(HEADERS, req_header) - @httpretty.activate def test_metadata(self): - httpretty.register_uri( - httpretty.GET, MD_URL_RE, - body=_new_request_callback()) + _set_mock_metadata() self.ds.get_data() shostname = GCE_META.get('instance/hostname').split('.')[0] @@ -107,18 +101,12 @@ class TestDataSourceGCE(test_helpers.HttprettyTestCase): self.assertEqual(GCE_META.get('instance/id'), self.ds.get_instance_id()) - self.assertEqual(GCE_META.get('instance/zone'), - self.ds.availability_zone) - self.assertEqual(GCE_META.get('instance/attributes/user-data'), self.ds.get_userdata_raw()) # test partial metadata (missing user-data in particular) - @httpretty.activate def test_metadata_partial(self): - httpretty.register_uri( - httpretty.GET, MD_URL_RE, - body=_new_request_callback(GCE_META_PARTIAL)) + _set_mock_metadata(GCE_META_PARTIAL) self.ds.get_data() self.assertEqual(GCE_META_PARTIAL.get('instance/id'), @@ -127,58 +115,47 @@ class TestDataSourceGCE(test_helpers.HttprettyTestCase): shostname = GCE_META_PARTIAL.get('instance/hostname').split('.')[0] self.assertEqual(shostname, self.ds.get_hostname()) - @httpretty.activate def test_metadata_encoding(self): - httpretty.register_uri( - httpretty.GET, MD_URL_RE, - body=_new_request_callback(GCE_META_ENCODING)) + _set_mock_metadata(GCE_META_ENCODING) self.ds.get_data() decoded = b64decode( GCE_META_ENCODING.get('instance/attributes/user-data')) self.assertEqual(decoded, self.ds.get_userdata_raw()) - @httpretty.activate def test_missing_required_keys_return_false(self): for required_key in ['instance/id', 'instance/zone', 'instance/hostname']: meta = GCE_META_PARTIAL.copy() del meta[required_key] - httpretty.register_uri(httpretty.GET, MD_URL_RE, - body=_new_request_callback(meta)) + _set_mock_metadata(meta) self.assertEqual(False, self.ds.get_data()) httpretty.reset() - @httpretty.activate def test_project_level_ssh_keys_are_used(self): - httpretty.register_uri(httpretty.GET, MD_URL_RE, - body=_new_request_callback()) + _set_mock_metadata() self.ds.get_data() # we expect a list of public ssh keys with user names stripped self.assertEqual(['ssh-rsa AA2..+aRD0fyVw== root@server'], self.ds.get_public_ssh_keys()) - @httpretty.activate def test_instance_level_ssh_keys_are_used(self): key_content = 'ssh-rsa JustAUser root@server' meta = GCE_META.copy() meta['instance/attributes/sshKeys'] = 'user:{0}'.format(key_content) - httpretty.register_uri(httpretty.GET, MD_URL_RE, - body=_new_request_callback(meta)) + _set_mock_metadata(meta) self.ds.get_data() self.assertIn(key_content, self.ds.get_public_ssh_keys()) - @httpretty.activate def test_instance_level_keys_replace_project_level_keys(self): key_content = 'ssh-rsa JustAUser root@server' meta = GCE_META.copy() meta['instance/attributes/sshKeys'] = 'user:{0}'.format(key_content) - httpretty.register_uri(httpretty.GET, MD_URL_RE, - body=_new_request_callback(meta)) + _set_mock_metadata(meta) self.ds.get_data() self.assertEqual([key_content], self.ds.get_public_ssh_keys()) |