Return a sensible value for DataSourceGCE.availability_zone.

LP: #1470880

3 files changed, 22 insertions, 35 deletions

diff --git a/ChangeLog b/ChangeLog
index 905b4fca..b6980eb0 100644
--- a/ChangeLog
+++ b/ChangeLog
@@ -51,6 +51,7 @@
  - cc_apt_configure: fix importing keys under python3 (LP: #1463373)
  - cc_growpart: fix specification of 'devices' list (LP: #1465436)
  - CloudStack: fix password setting on cloudstack > 4.5.1 (LP: #1464253)
+ - GCE: fix determination of availability zone (LP: #1470880)
 0.7.6:
  - open 0.7.6
  - Enable vendordata on CloudSigma datasource (LP: #1303986)
diff --git a/cloudinit/sources/DataSourceGCE.py b/cloudinit/sources/DataSourceGCE.py
index f4ed915d..1b28a68c 100644
--- a/cloudinit/sources/DataSourceGCE.py
+++ b/cloudinit/sources/DataSourceGCE.py
@@ -116,6 +116,10 @@ class DataSourceGCE(sources.DataSource):
             lines = self.metadata['public-keys'].splitlines()
             self.metadata['public-keys'] = [self._trim_key(k) for k in lines]
 
+        if self.metadata['availability-zone']:
+            self.metadata['availability-zone'] = self.metadata[
+                'availability-zone'].split('/')[-1]
+
         encoding = self.metadata.get('user-data-encoding')
         if encoding:
             if encoding == 'base64':
diff --git a/tests/unittests/test_datasource/test_gce.py b/tests/unittests/test_datasource/test_gce.py
index 1fb100f7..fa714070 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,52 @@ 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())
+
+    def test_only_last_part_of_zone_used_for_availability_zone(self):
+        _set_mock_metadata()
+        self.ds.get_data()
+        self.assertEqual('bar', self.ds.availability_zone)