Merge pull request #4 from zdc/T2117

Cloud-init: T2117: Updated to 20.1 version

1 files changed, 20 insertions, 2 deletions

diff --git a/tests/unittests/test_datasource/test_gce.py b/tests/unittests/test_datasource/test_gce.py
index 41176c6a..4afbccff 100644
--- a/tests/unittests/test_datasource/test_gce.py
+++ b/tests/unittests/test_datasource/test_gce.py
@@ -7,11 +7,11 @@
 import datetime
 import httpretty
 import json
-import mock
 import re
+from unittest import mock
+from urllib.parse import urlparse
 
 from base64 import b64encode, b64decode
-from six.moves.urllib_parse import urlparse
 
 from cloudinit import distros
 from cloudinit import helpers
@@ -55,6 +55,8 @@ GCE_USER_DATA_TEXT = {
 HEADERS = {'Metadata-Flavor': 'Google'}
 MD_URL_RE = re.compile(
     r'http://metadata.google.internal/computeMetadata/v1/.*')
+GUEST_ATTRIBUTES_URL = ('http://metadata.google.internal/computeMetadata/'
+                        'v1/instance/guest-attributes/hostkeys/')
 
 
 def _set_mock_metadata(gce_meta=None):
@@ -341,4 +343,20 @@ class TestDataSourceGCE(test_helpers.HttprettyTestCase):
             public_key_data, default_user='default')
         self.assertEqual(sorted(found), sorted(expected))
 
+    @mock.patch("cloudinit.url_helper.readurl")
+    def test_publish_host_keys(self, m_readurl):
+        hostkeys = [('ssh-rsa', 'asdfasdf'),
+                    ('ssh-ed25519', 'qwerqwer')]
+        readurl_expected_calls = [
+            mock.call(check_status=False, data=b'asdfasdf', headers=HEADERS,
+                      request_method='PUT',
+                      url='%s%s' % (GUEST_ATTRIBUTES_URL, 'ssh-rsa')),
+            mock.call(check_status=False, data=b'qwerqwer', headers=HEADERS,
+                      request_method='PUT',
+                      url='%s%s' % (GUEST_ATTRIBUTES_URL, 'ssh-ed25519')),
+        ]
+        self.ds.publish_host_keys(hostkeys)
+        m_readurl.assert_has_calls(readurl_expected_calls, any_order=True)
+
+
 # vi: ts=4 expandtab