Correct handling of single/multiple ssh keys

2 files changed, 31 insertions, 5 deletions

diff --git a/cloudinit/sources/DataSourceDigitalOcean.py b/cloudinit/sources/DataSourceDigitalOcean.py
index b7afca93..c59232ca 100644
--- a/cloudinit/sources/DataSourceDigitalOcean.py
+++ b/cloudinit/sources/DataSourceDigitalOcean.py
@@ -19,6 +19,7 @@ from cloudinit import util
 from cloudinit import sources
 from cloudinit import url_helper
 from cloudinit import ec2_utils
+from types import *
 import functools
 
 
@@ -71,7 +72,10 @@ class DataSourceDigitalOcean(sources.DataSource):
         return "\n".join(self.metadata['vendor-data'])
 
     def get_public_ssh_keys(self):
-        return self.metadata['public-keys'].splitlines()
+        if type(self.metadata['public-keys']) is StringType:
+           return [self.metadata['public-keys']]
+        else:
+           return self.metadata['public-keys']
 
     @property
     def availability_zone(self):
diff --git a/tests/unittests/test_datasource/test_digitalocean.py b/tests/unittests/test_datasource/test_digitalocean.py
index 559a4f9f..0997cf38 100644
--- a/tests/unittests/test_datasource/test_digitalocean.py
+++ b/tests/unittests/test_datasource/test_digitalocean.py
@@ -18,6 +18,7 @@
 import httpretty
 import re
 
+from types import *
 from urlparse import urlparse
 
 from cloudinit import settings
@@ -34,11 +35,15 @@ DO_INDEX = """id
            public-keys
            region"""
 
+DO_MULTIPLE_KEYS = """ssh-rsa AAAAB3NzaC1yc2EAAAA... neal@digitalocean.com
+                   ssh-rsa AAAAB3NzaC1yc2EAAAA... neal2@digitalocean.com"""
+DO_SINGLE_KEY = "ssh-rsa AAAAB3NzaC1yc2EAAAA... neal@digitalocean.com"
+
 DO_META = {
     '': DO_INDEX,
     'user-data': '#!/bin/bash\necho "user-data"',
     'vendor-data': '#!/bin/bash\necho "vendor-data"',
-    'public-keys': 'ssh-rsa AAAAB3NzaC1yc2EAAAA... neal@digitalocean.com',
+    'public-keys': DO_SINGLE_KEY,
     'region': 'nyc3',
     'id': '2000000',
     'hostname': 'cloudinit-test',
@@ -88,9 +93,6 @@ class TestDataSourceDigitalOcean(test_helpers.HttprettyTestCase):
         self.assertEqual(DO_META.get('vendor-data'),
                          self.ds.get_vendordata_raw())
 
-        self.assertEqual([DO_META.get('public-keys')],
-                         self.ds.get_public_ssh_keys())
-
         self.assertEqual(DO_META.get('region'),
                          self.ds.availability_zone)
 
@@ -102,3 +104,23 @@ class TestDataSourceDigitalOcean(test_helpers.HttprettyTestCase):
 
         self.assertEqual('http://mirrors.digitalocean.com/',
                          self.ds.get_package_mirror_info())
+
+        # Single key
+        self.assertEqual([DO_META.get('public-keys')],
+                         self.ds.get_public_ssh_keys())
+
+        self.assertIs(type(self.ds.get_public_ssh_keys()), ListType)
+
+    @httpretty.activate
+    def test_multiple_ssh_keys(self):
+        DO_META['public_keys'] = DO_MULTIPLE_KEYS
+        httpretty.register_uri(
+            httpretty.GET, MD_URL_RE,
+            body=_request_callback)
+        self.ds.get_data()
+
+        # Multiple keys
+        self.assertEqual(DO_META.get('public-keys').splitlines(),
+                         self.ds.get_public_ssh_keys())
+
+        self.assertIs(type(self.ds.get_public_ssh_keys()), ListType)