GCE: add some tests, some small fixes

2 files changed, 160 insertions, 28 deletions

diff --git a/cloudinit/sources/DataSourceGCE.py b/cloudinit/sources/DataSourceGCE.py
index 95a410ba..c993293f 100644
--- a/cloudinit/sources/DataSourceGCE.py
+++ b/cloudinit/sources/DataSourceGCE.py
@@ -31,7 +31,7 @@ REQUIRED_FIELDS = ('instance-id', 'availability-zone', 'local-hostname')
 class DataSourceGCE(sources.DataSource):
     def __init__(self, sys_cfg, distro, paths):
         sources.DataSource.__init__(self, sys_cfg, distro, paths)
-        self.metadata = {}
+        self.metadata = dict()
         self.ds_cfg = util.mergemanydict([
             util.get_cfg_by_path(sys_cfg, ["datasource", "GCE"], {}),
             BUILTIN_DS_CONFIG])
@@ -51,42 +51,60 @@ class DataSourceGCE(sources.DataSource):
         # GCE metadata server requires a custom header since v1
         headers = {'X-Google-Metadata-Request': True}
 
-        url_map = {
-            'instance-id': self.metadata_address + 'instance/id',
-            'availability-zone': self.metadata_address + 'instance/zone',
-            'public-keys': self.metadata_address + 'project/attributes/sshKeys',
-            'local-hostname': self.metadata_address + 'instance/hostname',
-            'user-data': self.metadata_address + 'instance/attributes/user-data',
-        }
+        # url_map: (our-key, path, required)
+        url_map = [
+            ('instance-id', 'instance/id', True),
+            ('availability-zone', 'instance/zone', True),
+            ('local-hostname', 'instance/hostname', True),
+            ('public-keys', 'project/attributes/sshKeys', False),
+            ('user-data', 'instance/attributes/user-data', False),
+        ]
 
         # if we cannot resolve the metadata server, then no point in trying
         if not util.is_resolvable(self.metadata_address):
             LOG.debug("%s is not resolvable", self.metadata_address)
             return False
 
-        for mkey in url_map.iterkeys():
+        # iterate over url_map keys to get metadata items
+        found = False
+        for (mkey, path, required) in url_map:
             try:
-                resp = url_helper.readurl(url=url_map[mkey], headers=headers,
-                                          retries=0)
-            except IOError:
-                return False
-            if resp.ok():
-                if mkey == 'public-keys':
-                    pub_keys = [self._trim_key(k) for k in resp.contents.splitlines()]
-                    self.metadata[mkey] = pub_keys
-                else:
+                resp = url_helper.readurl(url=self.metadata_address + path,
+                                          headers=headers)
+                if resp.code == 200:
+                    found = True
                     self.metadata[mkey] = resp.contents
-            else:
-                if mkey in REQUIRED_FIELDS:
-                    LOG.warn("required metadata '%s' not found in metadata",
-                        url_map[mkey])
+                else:
+                    if required:
+                        msg = "required url %s returned code %s. not GCE"
+                        if not found:
+                            LOG.debug(msg, path, resp.code)
+                        else:
+                            LOG.warn(msg, path, resp.code)
+                        return False
+                    else:
+                        self.metadata[mkey] = None
+            except url_helper.UrlError as e:
+                if required:
+                    msg = "required url %s raised exception %s. not GCE"
+                    if not found:
+                        LOG.debug(msg, path, e)
+                    else:
+                        LOG.warn(msg, path, e)
                     return False
-                        
+                msg = "Failed to get %s metadata item: %s."
+                if found:
+                    LOG.warn(msg, path, e)
+                else:
+                    LOG.debug(msg, path, e)
+
                 self.metadata[mkey] = None
-                return False
 
-        self.user_data_raw = self.metadata['user-data']
-        return True
+        if self.metadata['public-keys']:
+            lines = self.metadata['public-keys'].splitlines()
+            self.metadata['public-keys'] = [self._trim_key(k) for k in lines]
+
+        return found
 
     @property
     def launch_index(self):
@@ -99,12 +117,15 @@ class DataSourceGCE(sources.DataSource):
     def get_public_ssh_keys(self):
         return self.metadata['public-keys']
 
-    def get_hostname(self, fqdn=False):
+    def get_hostname(self, fqdn=False, _resolve_ip=False):
         return self.metadata['local-hostname']
 
+    def get_userdata_raw(self):
+        return self.metadata['user-data']
+
     @property
     def availability_zone(self):
-        return self.metadata['instance-zone']
+        return self.metadata['availability-zone']
 
 # Used to match classes to dependencies
 datasources = [
diff --git a/tests/unittests/test_datasource/test_gce.py b/tests/unittests/test_datasource/test_gce.py
new file mode 100644
index 00000000..d91bd531
--- /dev/null
+++ b/tests/unittests/test_datasource/test_gce.py
@@ -0,0 +1,111 @@
+#
+#    Copyright (C) 2014 Vaidas Jablonskis
+#
+#    Author: Vaidas Jablonskis <jablonskis@gmail.com>
+#
+#    This program is free software: you can redistribute it and/or modify
+#    it under the terms of the GNU General Public License version 3, as
+#    published by the Free Software Foundation.
+#
+#    This program is distributed in the hope that it will be useful,
+#    but WITHOUT ANY WARRANTY; without even the implied warranty of
+#    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+#    GNU General Public License for more details.
+#
+#    You should have received a copy of the GNU General Public License
+#    along with this program.  If not, see <http://www.gnu.org/licenses/>.
+
+import unittest
+import httpretty
+import re
+
+from urlparse import urlparse
+
+from cloudinit import settings
+from cloudinit import helpers
+from cloudinit.sources import DataSourceGCE
+
+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/attributes/user-data': '/bin/echo foo\n',
+}
+
+GCE_META_PARTIAL = {
+    'instance/id': '123',
+    'instance/hostname': 'server.project-name.local',
+}
+
+HEADERS = {'X-Google-Metadata-Request': 'True'}
+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, '')
+
+
+class TestDataSourceGCE(unittest.TestCase):
+
+    def setUp(self):
+        self.ds = DataSourceGCE.DataSourceGCE(
+            settings.CFG_BUILTIN, None,
+            helpers.Paths({}))
+
+    @httpretty.activate
+    def test_connection(self):
+        httpretty.register_uri(
+            httpretty.GET, MD_URL_RE,
+            body=_request_callback)
+
+        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=_request_callback)
+        self.ds.get_data()
+
+        self.assertEqual(GCE_META.get('instance/hostname'),
+                         self.ds.get_hostname())
+
+        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())
+
+        # 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())
+
+    # test partial metadata (missing user-data in particular)
+    @httpretty.activate
+    def test_metadata_partial(self):
+        httpretty.register_uri(
+            httpretty.GET, MD_URL_RE,
+            body=_request_callback)
+        self.ds.get_data()
+
+        self.assertEqual(GCE_META_PARTIAL.get('instance/id'),
+                         self.ds.get_instance_id())
+
+        self.assertEqual(GCE_META_PARTIAL.get('instance/hostname'),
+                         self.ds.get_hostname())