# Copyright 2013 Cloudbase Solutions Srl # # Licensed under the Apache License, Version 2.0 (the "License"); you may # not use this file except in compliance with the License. You may obtain # a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, WITHOUT # WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the # License for the specific language governing permissions and limitations # under the License. import os import pkgutil import tempfile import unittest try: import unittest.mock as mock except ImportError: import mock from cloudbaseinit.metadata.services import base as metadata_services_base from cloudbaseinit.plugins.common import base from cloudbaseinit.plugins.common import userdata from cloudbaseinit.tests.metadata import fake_json_response class FakeService(object): def __init__(self, user_data): self.user_data = user_data def get_user_data(self): return self.user_data.encode() def _create_tempfile(): fd, tmp = tempfile.mkstemp() os.close(fd) return tmp class UserDataPluginTest(unittest.TestCase): def setUp(self): self._userdata = userdata.UserDataPlugin() self.fake_data = fake_json_response.get_fake_metadata_json( '2013-04-04') @mock.patch('cloudbaseinit.plugins.common.userdata.UserDataPlugin' '._process_user_data') @mock.patch('cloudbaseinit.plugins.common.userdata.UserDataPlugin' '._check_gzip_compression') def _test_execute(self, mock_check_gzip_compression, mock_process_user_data, ret_val): mock_service = mock.MagicMock() mock_service.get_user_data.side_effect = [ret_val] response = self._userdata.execute(service=mock_service, shared_data=None) mock_service.get_user_data.assert_called_once_with() if ret_val is metadata_services_base.NotExistingMetadataException: self.assertEqual(response, (base.PLUGIN_EXECUTION_DONE, False)) elif ret_val is None: self.assertEqual(response, (base.PLUGIN_EXECUTION_DONE, False)) else: mock_check_gzip_compression.assert_called_once_with(ret_val) mock_process_user_data.assert_called_once_with( mock_check_gzip_compression.return_value) self.assertEqual(response, mock_process_user_data.return_value) def test_execute(self): self._test_execute(ret_val='fake_data') def test_execute_no_data(self): self._test_execute(ret_val=None) def test_execute_NotExistingMetadataException(self): self._test_execute( ret_val=metadata_services_base.NotExistingMetadataException) def test_execute_not_user_data(self): self._test_execute(ret_val=None) @mock.patch('io.BytesIO') @mock.patch('gzip.GzipFile') def test_check_gzip_compression(self, mock_GzipFile, mock_BytesIO): fake_userdata = b'\x1f\x8b' fake_userdata += self._userdata._GZIP_MAGIC_NUMBER response = self._userdata._check_gzip_compression(fake_userdata) mock_BytesIO.assert_called_once_with(fake_userdata) mock_GzipFile.assert_called_once_with( fileobj=mock_BytesIO.return_value, mode='rb') data = mock_GzipFile().__enter__().read.return_value self.assertEqual(data, response) @mock.patch('email.message_from_string') @mock.patch('cloudbaseinit.utils.encoding.get_as_string') def test_parse_mime(self, mock_get_as_string, mock_message_from_string): fake_user_data = 'fake data' response = self._userdata._parse_mime(user_data=fake_user_data) mock_get_as_string.assert_called_once_with(fake_user_data) mock_message_from_string.assert_called_once_with( mock_get_as_string.return_value) self.assertEqual(response, mock_message_from_string().walk()) @mock.patch('cloudbaseinit.plugins.common.userdataplugins.factory.' 'load_plugins') @mock.patch('cloudbaseinit.plugins.common.userdata.UserDataPlugin' '._parse_mime') @mock.patch('cloudbaseinit.plugins.common.userdata.UserDataPlugin' '._process_part') @mock.patch('cloudbaseinit.plugins.common.userdata.UserDataPlugin' '._end_part_process_event') @mock.patch('cloudbaseinit.plugins.common.userdata.UserDataPlugin' '._process_non_multi_part') def _test_process_user_data(self, mock_process_non_multi_part, mock_end_part_process_event, mock_process_part, mock_parse_mime, mock_load_plugins, user_data, reboot): mock_part = mock.MagicMock() mock_parse_mime.return_value = [mock_part] mock_process_part.return_value = (base.PLUGIN_EXECUTION_DONE, reboot) response = self._userdata._process_user_data(user_data=user_data) if user_data.startswith(b'Content-Type: multipart'): mock_load_plugins.assert_called_once_with() mock_parse_mime.assert_called_once_with(user_data) mock_process_part.assert_called_once_with(mock_part, mock_load_plugins(), {}) self.assertEqual((base.PLUGIN_EXECUTION_DONE, reboot), response) else: mock_process_non_multi_part.assert_called_once_with(user_data) self.assertEqual(mock_process_non_multi_part.return_value, response) def test_process_user_data_multipart_reboot_true(self): self._test_process_user_data(user_data=b'Content-Type: multipart', reboot=True) def test_process_user_data_multipart_reboot_false(self): self._test_process_user_data(user_data=b'Content-Type: multipart', reboot=False) def test_process_user_data_non_multipart(self): self._test_process_user_data(user_data=b'Content-Type: non-multipart', reboot=False) @mock.patch('cloudbaseinit.plugins.common.userdata.UserDataPlugin' '._add_part_handlers') @mock.patch('cloudbaseinit.plugins.common.execcmd' '.get_plugin_return_value') def _test_process_part(self, mock_get_plugin_return_value, mock_add_part_handlers, handler_func, user_data_plugin, content_type): mock_part = mock.MagicMock() mock_user_data_plugins = mock.MagicMock() mock_user_handlers = mock.MagicMock() mock_user_handlers.get.side_effect = [handler_func] mock_user_data_plugins.get.side_effect = [user_data_plugin] if content_type: _content_type = self._userdata._PART_HANDLER_CONTENT_TYPE mock_part.get_content_type.return_value = _content_type else: _content_type = 'other content type' mock_part.get_content_type.return_value = _content_type response = self._userdata._process_part( part=mock_part, user_data_plugins=mock_user_data_plugins, user_handlers=mock_user_handlers) mock_part.get_content_type.assert_called_once_with() mock_user_handlers.get.assert_called_once_with( _content_type) if handler_func and handler_func is Exception: self.assertEqual(2, mock_part.get_content_type.call_count) self.assertEqual(2, mock_part.get_filename.call_count) elif handler_func: handler_func.assert_called_once_with(None, _content_type, mock_part.get_filename(), mock_part.get_payload()) self.assertEqual(1, mock_part.get_content_type.call_count) self.assertEqual(2, mock_part.get_filename.call_count) else: mock_user_data_plugins.get.assert_called_once_with(_content_type) if user_data_plugin and content_type: user_data_plugin.process.assert_called_with(mock_part) mock_add_part_handlers.assert_called_with( mock_user_data_plugins, mock_user_handlers, user_data_plugin.process()) elif user_data_plugin and not content_type: mock_get_plugin_return_value.assert_called_once_with( user_data_plugin.process()) self.assertEqual(mock_get_plugin_return_value.return_value, response) def test_process_part(self): handler_func = mock.MagicMock() self._test_process_part(handler_func=handler_func, user_data_plugin=None, content_type=False) def test_process_part_no_handler_func(self): user_data_plugin = mock.MagicMock() self._test_process_part(handler_func=None, user_data_plugin=user_data_plugin, content_type=True) def test_process_part_not_content_type(self): user_data_plugin = mock.MagicMock() self._test_process_part(handler_func=None, user_data_plugin=user_data_plugin, content_type=False) @mock.patch('cloudbaseinit.plugins.common.userdata.UserDataPlugin' '._begin_part_process_event') def _test_add_part_handlers(self, mock_begin_part_process_event, ret_val): mock_user_data_plugins = mock.MagicMock(spec=dict) mock_new_user_handlers = mock.MagicMock(spec=dict) mock_user_handlers = mock.MagicMock(spec=dict) mock_handler_func = mock.MagicMock() mock_new_user_handlers.items.return_value = [('fake content type', mock_handler_func)] if ret_val: mock_user_data_plugins.get.return_value = mock_handler_func else: mock_user_data_plugins.get.return_value = None self._userdata._add_part_handlers( user_data_plugins=mock_user_data_plugins, user_handlers=mock_user_handlers, new_user_handlers=mock_new_user_handlers) mock_user_data_plugins.get.assert_called_with('fake content type') if ret_val is None: mock_user_handlers.__setitem__.assert_called_once_with( 'fake content type', mock_handler_func) mock_begin_part_process_event.assert_called_with(mock_handler_func) def test_add_part_handlers(self): self._test_add_part_handlers(ret_val=None) def test_add_part_handlers_skip_part_handler(self): mock_func = mock.MagicMock() self._test_add_part_handlers(ret_val=mock_func) def test_begin_part_process_event(self): mock_handler_func = mock.MagicMock() self._userdata._begin_part_process_event( handler_func=mock_handler_func) mock_handler_func.assert_called_once_with(None, "__begin__", None, None) def test_end_part_process_event(self): mock_handler_func = mock.MagicMock() self._userdata._end_part_process_event( handler_func=mock_handler_func) mock_handler_func.assert_called_once_with(None, "__end__", None, None) @mock.patch('cloudbaseinit.plugins.common.userdatautils' '.execute_user_data_script') @mock.patch('cloudbaseinit.plugins.common.execcmd' '.get_plugin_return_value') def test_process_non_multi_part(self, mock_get_plugin_return_value, mock_execute_user_data_script): user_data = b'fake' response = self._userdata._process_non_multi_part(user_data=user_data) mock_execute_user_data_script.assert_called_once_with(user_data) mock_get_plugin_return_value.assert_called_once_with( mock_execute_user_data_script()) self.assertEqual(mock_get_plugin_return_value.return_value, response) @mock.patch('cloudbaseinit.plugins.common.userdataplugins.factory.' 'load_plugins') @mock.patch('cloudbaseinit.plugins.common.execcmd' '.get_plugin_return_value') def test_process_non_multi_part_cloud_config( self, mock_get_plugin_return_value, mock_load_plugins): user_data = b'#cloud-config' mock_return_value = mock.sentinel.return_value mock_cloud_config_plugin = mock.Mock() mock_cloud_config_plugin.process.return_value = mock_return_value mock_get_plugin_return_value.return_value = mock_return_value mock_load_plugins.return_value = { 'text/cloud-config': mock_cloud_config_plugin} return_value = self._userdata._process_non_multi_part( user_data=user_data) mock_load_plugins.assert_called_once_with() (mock_cloud_config_plugin .process_non_multipart .assert_called_once_with(user_data)) self.assertEqual(mock_return_value, return_value) class TestCloudConfig(unittest.TestCase): @classmethod def setUpClass(cls): cls.plugin = userdata.UserDataPlugin() cls.userdata = pkgutil.get_data('cloudbaseinit.tests.resources', 'cloud_config_userdata').decode() def create_tempfiles(self, number): for _ in range(number): tmp = _create_tempfile() self.addCleanup(os.remove, tmp) yield tmp def test_cloud_config_multipart(self): b64, b64_binary, gz, gz_binary = list(self.create_tempfiles(4)) service = FakeService(self.userdata.format(b64=b64, b64_binary=b64_binary, gzip=gz, gzip_binary=gz_binary)) self.plugin.execute(service, {}) for path in (b64, b64_binary, gz, gz_binary): self.assertTrue(os.path.exists(path), "Path {} should exist.".format(path)) with open(path) as stream: self.assertEqual('42', stream.read())