Fixes based on comments (and a test improvement)

2 files changed, 21 insertions, 11 deletions

diff --git a/packages/python-google-compute-engine/google_compute_engine/instance_setup/instance_setup.py b/packages/python-google-compute-engine/google_compute_engine/instance_setup/instance_setup.py
index 3350c39..8e55a43 100755
--- a/packages/python-google-compute-engine/google_compute_engine/instance_setup/instance_setup.py
+++ b/packages/python-google-compute-engine/google_compute_engine/instance_setup/instance_setup.py
@@ -39,7 +39,8 @@ class PutRequest(urlrequest.Request):
 
 
 GUEST_ATTRIBUTES_URL = ('http://metadata.google.internal/computeMetadata/v1beta1/'
-                        'instance/guest-attributes/hostkeys/')
+                        'instance/guest-attributes')
+GA_NAMESPACE = 'hostkeys'
 
 
 class InstanceSetup(object):
@@ -133,6 +134,7 @@ class InstanceSetup(object):
     Args:
       key_type: string, the type of the SSH key.
       key_dest: string, a file location to store the SSH key.
+
     Returns:
       tuple, key_type and public key string.
     """
@@ -165,15 +167,15 @@ class InstanceSetup(object):
   def _WriteHostKeyToGuestAttributes(self, key_type, key_value):
     """Write a host key to guest attributes, ignoring errors."""
     headers = {'Metadata-Flavor': 'Google'}
-    url = GUEST_ATTRIBUTES_URL + key_type
+    url = '%s/%s/%s' % (GUEST_ATTRIBUTES_URL, GA_NAMESPACE, key_type)
     req = PutRequest(url, key_value, headers)
     try:
       response = urlrequest.urlopen(req)
       self.logger.debug(response)
-      self.logger.info('Wrote %s host key to guest attributes.' % key_type)
+      self.logger.info('Wrote %s host key to guest attributes.', key_type)
     except urlerror.HTTPError:
-      self.logger.info('Unable to write %s host key to guest attributes.'
-                       % key_type)
+      self.logger.info('Unable to write %s host key to guest attributes.',
+                       key_type)
 
   def _StartSshd(self):
     """Initialize the SSH daemon."""
diff --git a/packages/python-google-compute-engine/google_compute_engine/instance_setup/tests/instance_setup_test.py b/packages/python-google-compute-engine/google_compute_engine/instance_setup/tests/instance_setup_test.py
index ca238b1..a59ae0c 100644
--- a/packages/python-google-compute-engine/google_compute_engine/instance_setup/tests/instance_setup_test.py
+++ b/packages/python-google-compute-engine/google_compute_engine/instance_setup/tests/instance_setup_test.py
@@ -330,21 +330,29 @@ class InstanceSetupTest(unittest.TestCase):
     instance_setup.InstanceSetup._SetSshHostKeys(self.mock_setup)
     self.mock_instance_config.SetOption.assert_not_called()
 
-  @mock.patch('google_compute_engine.instance_setup.instance_setup.urlrequest.Request')
   @mock.patch('google_compute_engine.instance_setup.instance_setup.urlrequest.urlopen')
-  def testWriteHostKeyToGuestAttributes(self, mock_urlopen, mock_request):
+  @mock.patch('google_compute_engine.instance_setup.instance_setup.PutRequest')
+  def testWriteHostKeyToGuestAttributes(self, mock_put, mock_urlopen):
+    key_type = 'ssh-rsa'
+    key_value = 'asdfasdf'
+    expected_url = ('http://metadata.google.internal/computeMetadata/v1beta1/'
+                    'instance/guest-attributes/hostkeys/%s' % key_type)
+    headers = {'Metadata-Flavor': 'Google'}
+
     instance_setup.InstanceSetup._WriteHostKeyToGuestAttributes(
-        self.mock_setup, 'ssh-rsa', 'asdfasdf')
+        self.mock_setup, key_type, key_value)
     self.mock_logger.info.assert_called_with(
-        'Wrote ssh-rsa host key to guest attributes.')
+        'Wrote %s host key to guest attributes.', key_type)
+    mock_put.assert_called_with(expected_url, key_value, headers)
 
     mock_urlopen.side_effect = instance_setup.urlerror.HTTPError(
         'http://foo', 403, 'Forbidden', {}, None)
     instance_setup.InstanceSetup._WriteHostKeyToGuestAttributes(
-        self.mock_setup, 'ssh-rsa', 'asdfasdf')
+        self.mock_setup, key_type, key_value)
     self.mock_logger.info.assert_called_with(
-        'Unable to write ssh-rsa host key to guest attributes.')
+        'Unable to write %s host key to guest attributes.', key_type)
 
+  def testPutRequest(self):
     put_request = instance_setup.PutRequest('http://example.com/')
     self.assertEqual(put_request.get_method(), 'PUT')