diff options
Diffstat (limited to 'packages/python-google-compute-engine/google_compute_engine/instance_setup/tests/instance_setup_test.py')
| -rw-r--r-- | packages/python-google-compute-engine/google_compute_engine/instance_setup/tests/instance_setup_test.py | 72 |
1 files changed, 56 insertions, 16 deletions
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 97a4674..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 @@ -18,6 +18,7 @@ import subprocess from google_compute_engine.instance_setup import instance_setup +from google_compute_engine.test_compat import builtin from google_compute_engine.test_compat import mock from google_compute_engine.test_compat import unittest @@ -231,22 +232,33 @@ class InstanceSetupTest(unittest.TestCase): temp_dest = '/tmp/dest' mock_tempfile.return_value = mock_tempfile mock_tempfile.__enter__.return_value.name = temp_dest - - instance_setup.InstanceSetup._GenerateSshKey( - self.mock_setup, key_type, key_dest) - expected_calls = [ - mock.call.tempfile(prefix=key_type, delete=True), - mock.call.tempfile.__enter__(), - mock.call.tempfile.__exit__(None, None, None), - mock.call.logger.info(mock.ANY, key_dest), - mock.call.call( - ['ssh-keygen', '-t', key_type, '-f', temp_dest, '-N', '', '-q']), - mock.call.move(temp_dest, key_dest), - mock.call.move('%s.pub' % temp_dest, '%s.pub' % key_dest), - mock.call.permissions(key_dest, mode=0o600), - mock.call.permissions('%s.pub' % key_dest, mode=0o644), - ] - self.assertEqual(mocks.mock_calls, expected_calls) + mock_open = mock.mock_open() + key_file_contents = 'ssh-rsa asdfasdf' + expected_key_data = ('ssh-rsa', 'asdfasdf') + + with mock.patch('%s.open' % builtin, mock_open, create=False): + mock_open().read.return_value = key_file_contents + key_data = instance_setup.InstanceSetup._GenerateSshKey( + self.mock_setup, key_type, key_dest) + expected_calls = [ + mock.call.tempfile(prefix=key_type, delete=True), + mock.call.tempfile.__enter__(), + mock.call.tempfile.__exit__(None, None, None), + mock.call.logger.info(mock.ANY, key_dest), + mock.call.call( + ['ssh-keygen', '-t', key_type, '-f', temp_dest, '-N', '', '-q']), + mock.call.move(temp_dest, key_dest), + mock.call.move('%s.pub' % temp_dest, '%s.pub' % key_dest), + mock.call.permissions(key_dest, mode=0o600), + mock.call.permissions('%s.pub' % key_dest, mode=0o644), + ] + self.assertEqual(mocks.mock_calls, expected_calls) + self.assertEqual(key_data, expected_key_data) + + mock_open().read.return_value = '' + key_data = instance_setup.InstanceSetup._GenerateSshKey( + self.mock_setup, key_type, key_dest) + self.assertEqual(key_data, None) @mock.patch('google_compute_engine.instance_setup.instance_setup.subprocess.check_call') def testGenerateSshKeyProcessError(self, mock_call): @@ -318,6 +330,32 @@ 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.urlopen') + @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, key_type, key_value) + self.mock_logger.info.assert_called_with( + '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, key_type, key_value) + self.mock_logger.info.assert_called_with( + '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') + @mock.patch('google_compute_engine.instance_setup.instance_setup.os.listdir') def testSetSshHostKeysFirstBoot(self, mock_listdir): self.mock_instance_config.GetOptionString.return_value = None @@ -325,6 +363,7 @@ class InstanceSetupTest(unittest.TestCase): mock_instance_id.return_value = '123' self.mock_setup._GetInstanceId = mock_instance_id mock_generate_key = mock.Mock() + mock_generate_key.return_value = ('ssh-rsa', 'asdfasdf') self.mock_setup._GenerateSshKey = mock_generate_key mock_listdir.return_value = [ 'ssh_config', @@ -344,6 +383,7 @@ class InstanceSetupTest(unittest.TestCase): mock.call('ed25519', '/etc/ssh/ssh_host_ed25519_key'), mock.call('rsa', '/etc/ssh/ssh_host_rsa_key'), ] + self.assertEqual(sorted(mock_generate_key.mock_calls), expected_calls) self.mock_instance_config.SetOption.assert_called_once_with( 'Instance', 'instance_id', '123') |