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 | 61 |
1 files changed, 45 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..1a5c3fd 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,25 @@ 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() + + with mock.patch('%s.open' % builtin, mock_open, create=False): + mock_open().read.side_effect = 'ssh-rsa asdfasdf' + 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.patch('google_compute_engine.instance_setup.instance_setup.subprocess.check_call') def testGenerateSshKeyProcessError(self, mock_call): @@ -318,6 +322,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.PutRequest') + @mock.patch('google_compute_engine.instance_setup.instance_setup.urlrequest.urlopen') + def testWriteHostKeyToGuestAttributes(self, mock_urlopen, mock_put_request): + instance_setup.InstanceSetup._WriteHostKeyToGuestAttributes( + self.mock_setup, 'ssh-rsa', 'asdfasdf') + expected_calls = [ + mock.call('ssh-rsa', 'asdfasdf') + ] + self.mock_logger.info.assert_called_with( + 'Wrote ssh-rsa host key to guest attributes.') + + mock_urlopen.side_effect = instance_setup.urlerror.HTTPError( + 'http://foo', 403, 'Forbidden', {}, None) + instance_setup.InstanceSetup._WriteHostKeyToGuestAttributes( + self.mock_setup, 'ssh-rsa', 'asdfasdf') + expected_calls = [ + mock.call('ssh-rsa', 'asdfasdf') + ] + self.mock_logger.info.assert_called_with( + 'Unable to write ssh-rsa host key to guest attributes.') + + @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 +352,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 +372,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') |