1 files changed, 48 insertions, 6 deletions
diff --git a/tests/unittests/test_datasource/test_gce.py b/tests/unittests/test_datasource/test_gce.py
index 6dd4b5ed..1fb100f7 100644
--- a/tests/unittests/test_datasource/test_gce.py
+++ b/tests/unittests/test_datasource/test_gce.py
@@ -15,7 +15,6 @@
 #    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 httpretty
 import re
 
 from base64 import b64encode, b64decode
@@ -27,12 +26,14 @@ from cloudinit.sources import DataSourceGCE
 
 from .. import helpers as test_helpers
 
+httpretty = test_helpers.import_httpretty()
+
 GCE_META = {
     'instance/id': '123',
     'instance/zone': 'foo/bar',
     'project/attributes/sshKeys': 'user:ssh-rsa AA2..+aRD0fyVw== root@server',
     'instance/hostname': 'server.project-foo.local',
-    'instance/attributes/user-data': '/bin/echo foo\n',
+    'instance/attributes/user-data': b'/bin/echo foo\n',
 }
 
 GCE_META_PARTIAL = {
@@ -112,10 +113,6 @@ class TestDataSourceGCE(test_helpers.HttprettyTestCase):
         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):
@@ -140,3 +137,48 @@ class TestDataSourceGCE(test_helpers.HttprettyTestCase):
         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))
+            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())
+        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))
+        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))
+        self.ds.get_data()
+
+        self.assertEqual([key_content], self.ds.get_public_ssh_keys())