Add support for partition images in agent driver.

It also adds the ironic-lib in the requirements
list of the IPA package.

Partial-bug: 1526289
Depends-On: I22bc29a39bf5c35f3eecb6d4e51cebd6aee0ce19
Change-Id: I37908470484744bb720f741d378106d1cb1227a3

4 files changed, 343 insertions, 29 deletions

diff --git a/ironic_python_agent/extensions/standby.py b/ironic_python_agent/extensions/standby.py
index 238a524f..122b39ff 100644
--- a/ironic_python_agent/extensions/standby.py
+++ b/ironic_python_agent/extensions/standby.py
@@ -23,6 +23,7 @@ import time
 from oslo_concurrency import processutils
 from oslo_log import log
 
+from ironic_lib import disk_utils
 from ironic_python_agent import errors
 from ironic_python_agent.extensions import base
 from ironic_python_agent import hardware
@@ -46,10 +47,35 @@ def _path_to_script(script):
     return os.path.join(cwd, '..', script)
 
 
-def _write_image(image_info, device):
-    starttime = time.time()
-    image = _image_location(image_info)
+def _write_partition_image(image, image_info, device):
+    """Call disk_util to create partition and write the partition image."""
+    node_uuid = image_info['id']
+    preserve_ep = image_info['preserve_ephemeral']
+    configdrive = image_info['configdrive']
+    boot_option = image_info.get('boot_option', 'netboot')
+    boot_mode = image_info.get('deploy_boot_mode', 'bios')
+    image_mb = disk_utils.get_image_mb(image)
+    root_mb = image_info['root_mb']
+    if image_mb > int(root_mb):
+        msg = ('Root partition is too small for requested image. Image '
+               'virtual size: {0} MB, Root size: {1} MB').format(image_mb,
+                                                                 root_mb)
+        raise errors.InvalidCommandParamsError(msg)
+    try:
+        return disk_utils.work_on_disk(device, root_mb,
+                                       image_info['swap_mb'],
+                                       image_info['ephemeral_mb'],
+                                       image_info['ephemeral_format'],
+                                       image, node_uuid,
+                                       preserve_ephemeral=preserve_ep,
+                                       configdrive=configdrive,
+                                       boot_option=boot_option,
+                                       boot_mode=boot_mode)
+    except processutils.ProcessExecutionError as e:
+        raise errors.ImageWriteError(device, e.exit_code, e.stdout, e.stderr)
+
 
+def _write_whole_disk_image(image, image_info, device):
     script = _path_to_script('shell/write_image.sh')
     command = ['/bin/bash', script, image, device]
     LOG.info('Writing image with command: {0}'.format(' '.join(command)))
@@ -57,9 +83,20 @@ def _write_image(image_info, device):
         stdout, stderr = utils.execute(*command, check_exit_code=[0])
     except processutils.ProcessExecutionError as e:
         raise errors.ImageWriteError(device, e.exit_code, e.stdout, e.stderr)
+
+
+def _write_image(image_info, device):
+    starttime = time.time()
+    image = _image_location(image_info)
+    uuids = {}
+    if image_info.get('image_type') == 'partition':
+        uuids = _write_partition_image(image, image_info, device)
+    else:
+        _write_whole_disk_image(image, image_info, device)
     totaltime = time.time() - starttime
     LOG.info('Image {0} written to device {1} in {2} seconds'.format(
              image, device, totaltime))
+    return uuids
 
 
 def _configdrive_is_url(configdrive):
@@ -115,6 +152,27 @@ def _write_configdrive_to_partition(configdrive, device):
              totaltime))
 
 
+def _message_format(msg, image_info, device, partition_uuids):
+    """Helper method to get and populate different messages."""
+    message = None
+    result_msg = msg
+    if image_info.get('image_type') == 'partition':
+        root_uuid = partition_uuids.get('root uuid')
+        efi_system_partition_uuid = (
+            partition_uuids.get('efi system partition uuid'))
+        if image_info.get('deploy_boot_mode') == 'uefi':
+            result_msg = msg + 'root_uuid={2} efi_system_partition_uuid={3}'
+            message = result_msg.format(image_info['id'], device,
+                                        root_uuid,
+                                        efi_system_partition_uuid)
+        else:
+            result_msg = msg + 'root_uuid={2}'
+            message = result_msg.format(image_info['id'], device, root_uuid)
+    else:
+        message = result_msg.format(image_info['id'], device)
+    return message
+
+
 class ImageDownload(object):
     """Helper class that opens a HTTP connection to download an image.
 
@@ -219,10 +277,11 @@ class StandbyExtension(base.BaseAgentExtension):
         super(StandbyExtension, self).__init__(agent=agent)
 
         self.cached_image_id = None
+        self.partition_uuids = None
 
     def _cache_and_write_image(self, image_info, device):
         _download_image(image_info)
-        _write_image(image_info, device)
+        self.partition_uuids = _write_image(image_info, device)
         self.cached_image_id = image_info['id']
 
     def _stream_raw_image_onto_device(self, image_info, device):
@@ -249,17 +308,19 @@ class StandbyExtension(base.BaseAgentExtension):
         LOG.debug('Caching image %s', image_info['id'])
         device = hardware.dispatch_to_managers('get_os_install_device')
 
-        result_msg = 'image ({0}) already present on device {1}'
+        msg = 'image ({0}) already present on device {1} '
 
         if self.cached_image_id != image_info['id'] or force:
             LOG.debug('Already had %s cached, overwriting',
                       self.cached_image_id)
             self._cache_and_write_image(image_info, device)
-            result_msg = 'image ({0}) cached to device {1}'
+            msg = 'image ({0}) cached to device {1} '
 
-        msg = result_msg.format(image_info['id'], device)
-        LOG.info(msg)
-        return msg
+        result_msg = _message_format(msg, image_info, device,
+                                     self.partition_uuids)
+
+        LOG.info(result_msg)
+        return result_msg
 
     @base.async_command('prepare_image', _validate_image_info)
     def prepare_image(self,
@@ -277,18 +338,23 @@ class StandbyExtension(base.BaseAgentExtension):
                 LOG.debug('Already had %s cached, overwriting',
                           self.cached_image_id)
 
-            if stream_raw_images and disk_format == 'raw':
+            if (stream_raw_images and disk_format == 'raw' and
+                image_info.get('image_type') != 'partition'):
                 self._stream_raw_image_onto_device(image_info, device)
             else:
                 self._cache_and_write_image(image_info, device)
 
-        if configdrive is not None:
-            _write_configdrive_to_partition(configdrive, device)
-
-        msg = ('image ({0}) written to device {1}'.format(
-            image_info['id'], device))
-        LOG.info(msg)
-        return msg
+        # the configdrive creation is taken care by ironic-lib's
+        # work_on_disk().
+        if image_info.get('image_type') != 'partition':
+            if configdrive is not None:
+                _write_configdrive_to_partition(configdrive, device)
+
+        msg = 'image ({0}) written to device {1} '
+        result_msg = _message_format(msg, image_info, device,
+                                     self.partition_uuids)
+        LOG.info(result_msg)
+        return result_msg
 
     def _run_shutdown_script(self, parameter):
         script = _path_to_script('shell/shutdown.sh')
diff --git a/ironic_python_agent/tests/unit/extensions/test_standby.py b/ironic_python_agent/tests/unit/extensions/test_standby.py
index 75d2ce96..2e98cc45 100644
--- a/ironic_python_agent/tests/unit/extensions/test_standby.py
+++ b/ironic_python_agent/tests/unit/extensions/test_standby.py
@@ -32,6 +32,24 @@ def _build_fake_image_info():
     }
 
 
+def _build_fake_partition_image_info():
+    return {
+        'id': 'fake_id',
+        'urls': [
+            'http://example.org',
+        ],
+        'checksum': 'abc123',
+        'root_mb': '10',
+        'swap_mb': '10',
+        'ephemeral_mb': '10',
+        'ephemeral_format': 'abc',
+        'preserve_ephemeral': 'False',
+        'configdrive': 'configdrive',
+        'image_type': 'partition',
+        'boot_option': 'netboot',
+        'deploy_boot_mode': 'bios'}
+
+
 class TestStandbyExtension(test_base.BaseTestCase):
     def setUp(self):
         super(TestStandbyExtension, self).setUp()
@@ -115,6 +133,104 @@ class TestStandbyExtension(test_base.BaseTestCase):
 
         execute_mock.assert_called_once_with(*command, check_exit_code=[0])
 
+    @mock.patch('six.moves.builtins.open', autospec=True)
+    @mock.patch('ironic_python_agent.utils.execute', autospec=True)
+    @mock.patch('ironic_lib.disk_utils.get_image_mb', autospec=True)
+    @mock.patch('ironic_lib.disk_utils.work_on_disk', autospec=True)
+    def test_write_partition_image_exception(self, work_on_disk_mock,
+                                             image_mb_mock,
+                                             execute_mock, open_mock):
+        image_info = _build_fake_partition_image_info()
+        device = '/dev/sda'
+        root_mb = image_info['root_mb']
+        swap_mb = image_info['swap_mb']
+        ephemeral_mb = image_info['ephemeral_mb']
+        ephemeral_format = image_info['ephemeral_format']
+        node_uuid = image_info['id']
+        pr_ep = image_info['preserve_ephemeral']
+        configdrive = image_info['configdrive']
+        boot_mode = image_info['deploy_boot_mode']
+        boot_option = image_info['boot_option']
+
+        image_path = standby._image_location(image_info)
+
+        image_mb_mock.return_value = 1
+        exc = errors.ImageWriteError
+        Exception_returned = processutils.ProcessExecutionError
+        work_on_disk_mock.side_effect = Exception_returned
+
+        self.assertRaises(exc, standby._write_image, image_info,
+                          device)
+        image_mb_mock.assert_called_once_with(image_path)
+        work_on_disk_mock.assert_called_once_with(device, root_mb, swap_mb,
+                                                  ephemeral_mb,
+                                                  ephemeral_format,
+                                                  image_path,
+                                                  node_uuid,
+                                                  configdrive=configdrive,
+                                                  preserve_ephemeral=pr_ep,
+                                                  boot_mode=boot_mode,
+                                                  boot_option=boot_option)
+
+    @mock.patch('six.moves.builtins.open', autospec=True)
+    @mock.patch('ironic_python_agent.utils.execute', autospec=True)
+    @mock.patch('ironic_lib.disk_utils.get_image_mb', autospec=True)
+    @mock.patch('ironic_lib.disk_utils.work_on_disk', autospec=True)
+    def test_write_partition_image_exception_image_mb(self,
+                                                      work_on_disk_mock,
+                                                      image_mb_mock,
+                                                      execute_mock,
+                                                      open_mock):
+        image_info = _build_fake_partition_image_info()
+        device = '/dev/sda'
+        image_path = standby._image_location(image_info)
+
+        image_mb_mock.return_value = 20
+        exc = errors.InvalidCommandParamsError
+
+        self.assertRaises(exc, standby._write_image, image_info,
+                          device)
+        image_mb_mock.assert_called_once_with(image_path)
+        self.assertFalse(work_on_disk_mock.called)
+
+    @mock.patch('six.moves.builtins.open', autospec=True)
+    @mock.patch('ironic_python_agent.utils.execute', autospec=True)
+    @mock.patch('ironic_lib.disk_utils.work_on_disk', autospec=True)
+    @mock.patch('ironic_lib.disk_utils.get_image_mb', autospec=True)
+    def test_write_partition_image(self, image_mb_mock, work_on_disk_mock,
+                                   execute_mock, open_mock):
+        image_info = _build_fake_partition_image_info()
+        device = '/dev/sda'
+        root_mb = image_info['root_mb']
+        swap_mb = image_info['swap_mb']
+        ephemeral_mb = image_info['ephemeral_mb']
+        ephemeral_format = image_info['ephemeral_format']
+        node_uuid = image_info['id']
+        pr_ep = image_info['preserve_ephemeral']
+        configdrive = image_info['configdrive']
+        boot_mode = image_info['deploy_boot_mode']
+        boot_option = image_info['boot_option']
+
+        image_path = standby._image_location(image_info)
+        uuids = {'root uuid': 'root_uuid'}
+        expected_uuid = {'root uuid': 'root_uuid'}
+        image_mb_mock.return_value = 1
+        work_on_disk_mock.return_value = uuids
+
+        standby._write_image(image_info, device)
+        image_mb_mock.assert_called_once_with(image_path)
+        work_on_disk_mock.assert_called_once_with(device, root_mb, swap_mb,
+                                                  ephemeral_mb,
+                                                  ephemeral_format,
+                                                  image_path,
+                                                  node_uuid,
+                                                  configdrive=configdrive,
+                                                  preserve_ephemeral=pr_ep,
+                                                  boot_mode=boot_mode,
+                                                  boot_option=boot_option)
+
+        self.assertEqual(expected_uuid, work_on_disk_mock.return_value)
+
     def test_configdrive_is_url(self):
         self.assertTrue(standby._configdrive_is_url('http://some/url'))
         self.assertTrue(standby._configdrive_is_url('https://some/url'))
@@ -303,8 +419,34 @@ class TestStandbyExtension(test_base.BaseTestCase):
                          self.agent_extension.cached_image_id)
         self.assertEqual('SUCCEEDED', async_result.command_status)
         self.assertTrue('result' in async_result.command_result.keys())
-        cmd_result = ('cache_image: image ({0}) cached to device {1}'
-                      ).format(image_info['id'], 'manager')
+        cmd_result = ('cache_image: image ({0}) cached to device '
+                      '{1} ').format(image_info['id'], 'manager')
+        self.assertEqual(cmd_result, async_result.command_result['result'])
+
+    @mock.patch('ironic_python_agent.hardware.dispatch_to_managers',
+                autospec=True)
+    @mock.patch('ironic_python_agent.extensions.standby._write_image',
+                autospec=True)
+    @mock.patch('ironic_python_agent.extensions.standby._download_image',
+                autospec=True)
+    def test_cache_partition_image(self, download_mock, write_mock,
+                                   dispatch_mock):
+        image_info = _build_fake_partition_image_info()
+        download_mock.return_value = None
+        write_mock.return_value = {'root uuid': 'root_uuid'}
+        dispatch_mock.return_value = 'manager'
+        async_result = self.agent_extension.cache_image(image_info=image_info)
+        async_result.join()
+        download_mock.assert_called_once_with(image_info)
+        write_mock.assert_called_once_with(image_info, 'manager')
+        dispatch_mock.assert_called_once_with('get_os_install_device')
+        self.assertEqual(image_info['id'],
+                         self.agent_extension.cached_image_id)
+        self.assertEqual('SUCCEEDED', async_result.command_status)
+        self.assertTrue('result' in async_result.command_result.keys())
+        cmd_result = ('cache_image: image ({0}) cached to device {1} '
+                      'root_uuid={2}').format(image_info['id'], 'manager',
+                                              'root_uuid')
         self.assertEqual(cmd_result, async_result.command_result['result'])
 
     @mock.patch('ironic_python_agent.hardware.dispatch_to_managers',
@@ -331,8 +473,8 @@ class TestStandbyExtension(test_base.BaseTestCase):
                          self.agent_extension.cached_image_id)
         self.assertEqual('SUCCEEDED', async_result.command_status)
         self.assertTrue('result' in async_result.command_result.keys())
-        cmd_result = ('cache_image: image ({0}) cached to device {1}'
-                      ).format(image_info['id'], 'manager')
+        cmd_result = ('cache_image: image ({0}) cached to device '
+                      '{1} ').format(image_info['id'], 'manager')
         self.assertEqual(cmd_result, async_result.command_result['result'])
 
     @mock.patch('ironic_python_agent.hardware.dispatch_to_managers',
@@ -357,8 +499,8 @@ class TestStandbyExtension(test_base.BaseTestCase):
                          self.agent_extension.cached_image_id)
         self.assertEqual('SUCCEEDED', async_result.command_status)
         self.assertTrue('result' in async_result.command_result.keys())
-        cmd_result = ('cache_image: image ({0}) already present on device {1}'
-                      ).format(image_info['id'], 'manager')
+        cmd_result = ('cache_image: image ({0}) already present on device '
+                      '{1} ').format(image_info['id'], 'manager')
         self.assertEqual(cmd_result, async_result.command_result['result'])
 
     @mock.patch(('ironic_python_agent.extensions.standby.'
@@ -399,8 +541,8 @@ class TestStandbyExtension(test_base.BaseTestCase):
 
         self.assertEqual('SUCCEEDED', async_result.command_status)
         self.assertTrue('result' in async_result.command_result.keys())
-        cmd_result = ('prepare_image: image ({0}) written to device {1}'
-                      ).format(image_info['id'], 'manager')
+        cmd_result = ('prepare_image: image ({0}) written to device '
+                      '{1} ').format(image_info['id'], 'manager')
         self.assertEqual(cmd_result, async_result.command_result['result'])
 
         download_mock.reset_mock()
@@ -420,8 +562,71 @@ class TestStandbyExtension(test_base.BaseTestCase):
 
         self.assertEqual('SUCCEEDED', async_result.command_status)
         self.assertTrue('result' in async_result.command_result.keys())
-        cmd_result = ('prepare_image: image ({0}) written to device {1}'
-                      ).format(image_info['id'], 'manager')
+        cmd_result = ('prepare_image: image ({0}) written to device '
+                      '{1} ').format(image_info['id'], 'manager')
+        self.assertEqual(cmd_result, async_result.command_result['result'])
+
+    @mock.patch(('ironic_python_agent.extensions.standby.'
+                 '_write_configdrive_to_partition'),
+                autospec=True)
+    @mock.patch('ironic_python_agent.hardware.dispatch_to_managers',
+                autospec=True)
+    @mock.patch('ironic_python_agent.extensions.standby._write_image',
+                autospec=True)
+    @mock.patch('ironic_python_agent.extensions.standby._download_image',
+                autospec=True)
+    @mock.patch('ironic_python_agent.extensions.standby._configdrive_location',
+                autospec=True)
+    def test_prepare_partition_image(self,
+                                     location_mock,
+                                     download_mock,
+                                     write_mock,
+                                     dispatch_mock,
+                                     configdrive_copy_mock):
+        image_info = _build_fake_partition_image_info()
+        location_mock.return_value = '/tmp/configdrive'
+        download_mock.return_value = None
+        write_mock.return_value = {'root uuid': 'root_uuid'}
+        dispatch_mock.return_value = 'manager'
+        configdrive_copy_mock.return_value = None
+
+        async_result = self.agent_extension.prepare_image(
+            image_info=image_info,
+            configdrive='configdrive_data'
+        )
+        async_result.join()
+
+        download_mock.assert_called_once_with(image_info)
+        write_mock.assert_called_once_with(image_info, 'manager')
+        dispatch_mock.assert_called_once_with('get_os_install_device')
+        self.assertFalse(configdrive_copy_mock.called)
+
+        self.assertEqual('SUCCEEDED', async_result.command_status)
+        self.assertTrue('result' in async_result.command_result.keys())
+        cmd_result = ('prepare_image: image ({0}) written to device {1} '
+                      'root_uuid={2}').format(
+            image_info['id'], 'manager', 'root_uuid')
+        self.assertEqual(cmd_result, async_result.command_result['result'])
+
+        download_mock.reset_mock()
+        write_mock.reset_mock()
+        configdrive_copy_mock.reset_mock()
+        # image is now cached, make sure download/write doesn't happen
+        async_result = self.agent_extension.prepare_image(
+            image_info=image_info,
+            configdrive='configdrive_data'
+        )
+        async_result.join()
+
+        self.assertEqual(0, download_mock.call_count)
+        self.assertEqual(0, write_mock.call_count)
+        self.assertFalse(configdrive_copy_mock.called)
+
+        self.assertEqual('SUCCEEDED', async_result.command_status)
+        self.assertTrue('result' in async_result.command_result.keys())
+        cmd_result = ('prepare_image: image ({0}) written to device {1} '
+                      'root_uuid={2}').format(
+            image_info['id'], 'manager', 'root_uuid')
         self.assertEqual(cmd_result, async_result.command_result['result'])
 
     @mock.patch(('ironic_python_agent.extensions.standby.'
@@ -457,8 +662,8 @@ class TestStandbyExtension(test_base.BaseTestCase):
         self.assertEqual(0, configdrive_copy_mock.call_count)
         self.assertEqual('SUCCEEDED', async_result.command_status)
         self.assertTrue('result' in async_result.command_result.keys())
-        cmd_result = ('prepare_image: image ({0}) written to device {1}'
-                      ).format(image_info['id'], 'manager')
+        cmd_result = ('prepare_image: image ({0}) written to device '
+                      '{1} ').format(image_info['id'], 'manager')
         self.assertEqual(cmd_result, async_result.command_result['result'])
 
     @mock.patch(('ironic_python_agent.extensions.standby.'
@@ -611,6 +816,43 @@ class TestStandbyExtension(test_base.BaseTestCase):
         # Assert write was only called once and failed!
         file_mock.write.assert_called_once_with('some')
 
+    def test__message_format_whole_disk(self):
+        image_info = _build_fake_image_info()
+        msg = 'image ({0}) already present on device {1}'
+        device = '/dev/fake'
+        partition_uuids = {}
+        result_msg = standby._message_format(msg, image_info,
+                                             device, partition_uuids)
+        expected_msg = ('image (fake_id) already present on device '
+                        '/dev/fake')
+        self.assertEqual(expected_msg, result_msg)
+
+    def test__message_format_partition_bios(self):
+        image_info = _build_fake_partition_image_info()
+        msg = ('image ({0}) already present on device {1} ')
+        device = '/dev/fake'
+        partition_uuids = {'root uuid': 'root_uuid',
+                           'efi system partition uuid': None}
+        result_msg = standby._message_format(msg, image_info,
+                                             device, partition_uuids)
+        expected_msg = ('image (fake_id) already present on device '
+                        '/dev/fake root_uuid=root_uuid')
+        self.assertEqual(expected_msg, result_msg)
+
+    def test__message_format_partition_uefi(self):
+        image_info = _build_fake_partition_image_info()
+        image_info['deploy_boot_mode'] = 'uefi'
+        msg = ('image ({0}) already present on device {1} ')
+        device = '/dev/fake'
+        partition_uuids = {'root uuid': 'root_uuid',
+                           'efi system partition uuid': 'efi_id'}
+        result_msg = standby._message_format(msg, image_info,
+                                             device, partition_uuids)
+        expected_msg = ('image (fake_id) already present on device '
+                        '/dev/fake root_uuid=root_uuid '
+                        'efi_system_partition_uuid=efi_id')
+        self.assertEqual(expected_msg, result_msg)
+
 
 class TestImageDownload(test_base.BaseTestCase):
 
diff --git a/releasenotes/notes/agent_partition_image-91941adc6683c673.yaml b/releasenotes/notes/agent_partition_image-91941adc6683c673.yaml
new file mode 100644
index 00000000..fade2922
--- /dev/null
+++ b/releasenotes/notes/agent_partition_image-91941adc6683c673.yaml
@@ -0,0 +1,5 @@
+---
+features:
+  - Add support for partition images in IPA.
+    This commit adds the ironic-lib as the
+    requirement for the IPA package.
diff --git a/requirements.txt b/requirements.txt
index 0752ba6b..b17c8d09 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -22,3 +22,4 @@ rtslib-fb>=2.1.41 # Apache-2.0
 six>=1.9.0 # MIT
 stevedore>=1.5.0 # Apache-2.0
 WSME>=0.8 # MIT
+ironic-lib>=1.1.0 # Apache-2.0