Fix +tests for GCE datasource not handling per-instance SSH keys

1 files changed, 50 insertions, 42 deletions

diff --git a/cloudinit/sources/DataSourceGCE.py b/cloudinit/sources/DataSourceGCE.py
index 608c07f1..f4ed915d 100644
--- a/cloudinit/sources/DataSourceGCE.py
+++ b/cloudinit/sources/DataSourceGCE.py
@@ -30,6 +30,31 @@ BUILTIN_DS_CONFIG = {
 REQUIRED_FIELDS = ('instance-id', 'availability-zone', 'local-hostname')
 
 
+class GoogleMetadataFetcher(object):
+    headers = {'X-Google-Metadata-Request': True}
+
+    def __init__(self, metadata_address):
+        self.metadata_address = metadata_address
+
+    def get_value(self, path, is_text):
+        value = None
+        try:
+            resp = url_helper.readurl(url=self.metadata_address + path,
+                                      headers=self.headers)
+        except url_helper.UrlError as exc:
+            msg = "url %s raised exception %s"
+            LOG.debug(msg, path, exc)
+        else:
+            if resp.code == 200:
+                if is_text:
+                    value = util.decode_binary(resp.contents)
+                else:
+                    value = resp.contents
+            else:
+                LOG.debug("url %s returned code %s", path, resp.code)
+        return value
+
+
 class DataSourceGCE(sources.DataSource):
     def __init__(self, sys_cfg, distro, paths):
         sources.DataSource.__init__(self, sys_cfg, distro, paths)
@@ -50,17 +75,15 @@ class DataSourceGCE(sources.DataSource):
             return public_key
 
     def get_data(self):
-        # GCE metadata server requires a custom header since v1
-        headers = {'X-Google-Metadata-Request': True}
-
         # url_map: (our-key, path, required, is_text)
         url_map = [
-            ('instance-id', 'instance/id', True, True),
-            ('availability-zone', 'instance/zone', True, True),
-            ('local-hostname', 'instance/hostname', True, True),
-            ('public-keys', 'project/attributes/sshKeys', False, True),
-            ('user-data', 'instance/attributes/user-data', False, False),
-            ('user-data-encoding', 'instance/attributes/user-data-encoding',
+            ('instance-id', ('instance/id',), True, True),
+            ('availability-zone', ('instance/zone',), True, True),
+            ('local-hostname', ('instance/hostname',), True, True),
+            ('public-keys', ('project/attributes/sshKeys',
+                             'instance/attributes/sshKeys'), False, True),
+            ('user-data', ('instance/attributes/user-data',), False, False),
+            ('user-data-encoding', ('instance/attributes/user-data-encoding',),
              False, True),
         ]
 
@@ -69,40 +92,25 @@ class DataSourceGCE(sources.DataSource):
             LOG.debug("%s is not resolvable", self.metadata_address)
             return False
 
+        metadata_fetcher = GoogleMetadataFetcher(self.metadata_address)
         # iterate over url_map keys to get metadata items
-        found = False
-        for (mkey, path, required, is_text) in url_map:
-            try:
-                resp = url_helper.readurl(url=self.metadata_address + path,
-                                          headers=headers)
-                if resp.code == 200:
-                    found = True
-                    if is_text:
-                        self.metadata[mkey] = util.decode_binary(resp.contents)
-                    else:
-                        self.metadata[mkey] = resp.contents
+        running_on_gce = False
+        for (mkey, paths, required, is_text) in url_map:
+            value = None
+            for path in paths:
+                new_value = metadata_fetcher.get_value(path, is_text)
+                if new_value is not None:
+                    value = new_value
+            if value:
+                running_on_gce = True
+            if required and value is None:
+                msg = "required key %s returned nothing. not GCE"
+                if not running_on_gce:
+                    LOG.debug(msg, 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."
-                LOG.debug(msg, path, e)
-
-                self.metadata[mkey] = None
+                    LOG.warn(msg, mkey)
+                return False
+            self.metadata[mkey] = value
 
         if self.metadata['public-keys']:
             lines = self.metadata['public-keys'].splitlines()
@@ -116,7 +124,7 @@ class DataSourceGCE(sources.DataSource):
             else:
                 LOG.warn('unknown user-data-encoding: %s, ignoring', encoding)
 
-        return found
+        return running_on_gce
 
     @property
     def launch_index(self):