diff options
Diffstat (limited to 'tests/unittests/test_datasource')
| -rw-r--r-- | tests/unittests/test_datasource/test_gce.py | 60 | 
1 files changed, 44 insertions, 16 deletions
| diff --git a/tests/unittests/test_datasource/test_gce.py b/tests/unittests/test_datasource/test_gce.py index 842a72ba..e2449ba8 100644 --- a/tests/unittests/test_datasource/test_gce.py +++ b/tests/unittests/test_datasource/test_gce.py @@ -18,6 +18,7 @@  import httpretty  import re +from base64 import b64encode, b64decode  from urlparse import urlparse  from cloudinit import settings @@ -30,13 +31,23 @@ GCE_META = {      'instance/id': '123',      'instance/zone': 'foo/bar',      'project/attributes/sshKeys': 'user:ssh-rsa AA2..+aRD0fyVw== root@server', -    'instance/hostname': 'server.project-name.local', +    'instance/hostname': 'server.project-foo.local',      'instance/attributes/user-data': '/bin/echo foo\n', +    'instance/attributes/user-data-encoding': '',  }  GCE_META_PARTIAL = { -    'instance/id': '123', -    'instance/hostname': 'server.project-name.local', +    'instance/id': '1234', +    'instance/hostname': 'server.project-bar.local', +    'instance/zone': 'bar/baz', +} + +GCE_META_ENCODING = { +    'instance/id': '12345', +    'instance/hostname': 'server.project-baz.local', +    'instance/zone': 'baz/bang', +    'instance/attributes/user-data': b64encode('/bin/echo baz\n'), +    'instance/attributes/user-data-encoding': 'base64',  }  HEADERS = {'X-Google-Metadata-Request': 'True'} @@ -44,16 +55,22 @@ MD_URL_RE = re.compile(      r'http://metadata.google.internal./computeMetadata/v1/.*') -def _request_callback(method, uri, headers): -    url_path = urlparse(uri).path -    if url_path.startswith('/computeMetadata/v1/'): -        path = url_path.split('/computeMetadata/v1/')[1:][0] -    else: -        path = None -    if path in GCE_META: -        return (200, headers, GCE_META.get(path)) -    else: -        return (404, headers, '') +def _new_request_callback(gce_meta=None): +    if not gce_meta: +        gce_meta = GCE_META + +    def _request_callback(method, uri, headers): +        url_path = urlparse(uri).path +        if url_path.startswith('/computeMetadata/v1/'): +            path = url_path.split('/computeMetadata/v1/')[1:][0] +        else: +            path = None +        if path in gce_meta: +            return (200, headers, gce_meta.get(path)) +        else: +            return (404, headers, '') + +    return _request_callback  class TestDataSourceGCE(test_helpers.HttprettyTestCase): @@ -68,7 +85,7 @@ class TestDataSourceGCE(test_helpers.HttprettyTestCase):      def test_connection(self):          httpretty.register_uri(              httpretty.GET, MD_URL_RE, -            body=_request_callback) +            body=_new_request_callback())          success = self.ds.get_data()          self.assertTrue(success) @@ -80,7 +97,7 @@ class TestDataSourceGCE(test_helpers.HttprettyTestCase):      def test_metadata(self):          httpretty.register_uri(              httpretty.GET, MD_URL_RE, -            body=_request_callback) +            body=_new_request_callback())          self.ds.get_data()          self.assertEqual(GCE_META.get('instance/hostname'), @@ -104,7 +121,7 @@ class TestDataSourceGCE(test_helpers.HttprettyTestCase):      def test_metadata_partial(self):          httpretty.register_uri(              httpretty.GET, MD_URL_RE, -            body=_request_callback) +            body=_new_request_callback(GCE_META_PARTIAL))          self.ds.get_data()          self.assertEqual(GCE_META_PARTIAL.get('instance/id'), @@ -112,3 +129,14 @@ class TestDataSourceGCE(test_helpers.HttprettyTestCase):          self.assertEqual(GCE_META_PARTIAL.get('instance/hostname'),                           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)) +        self.ds.get_data() + +        decoded = b64decode( +            GCE_META_ENCODING.get('instance/attributes/user-data')) +        self.assertEqual(decoded, self.ds.get_userdata_raw()) | 
