Combine "can this be" and "make this a" boot disk logic for UEFI

This one does change behaviour in that if it has to resize a partition
to make room for the ESP it puts it before the partition, not after.
This lets us put an assertion back in an API test.

subiquity/common/filesystem/boot.py

@@ -97,6 +97,14 @@ class SetAttrPlan:
         setattr(self.device, self.attr, self.val)
 
 
+@attr.s(auto_attribs=True)
+class MountBootEfiPlan:
+    part: object
+
+    def apply(self, manipulator):
+        manipulator._mount_esp(self.part)
+
+
 @attr.s(auto_attribs=True)
 class MultiStepPlan:
     plans: list
@@ -148,6 +156,56 @@ def get_boot_device_plan_bios(device):
             ])
 
 
+def get_boot_device_plan_uefi(device):
+    if device._has_preexisting_partition():
+        for part in device.partitions():
+            if is_esp(part):
+                plans = [SetAttrPlan(part, 'grub_device', True)]
+                if device._m._mount_for_path('/boot/efi') is None:
+                    plans.append(MountBootEfiPlan(part))
+                return MultiStepPlan(plans=plans)
+        return None
+
+    size = sizes.get_efi_size(device)
+
+    create_part_plan = CreatePartPlan(
+        device=device,
+        offset=None,
+        spec=dict(size=size, fstype='fat32', mount=None),
+        args=dict(flag='boot', grub_device=True))
+    if device._m._mount_for_path("/boot/efi") is None:
+        create_part_plan.spec['mount'] = '/boot/efi'
+
+    partitions = device.partitions()
+
+    if gaps.largest_gap_size(device) >= size:
+        create_part_plan.offset = gaps.largest_gap(device).offset
+        return create_part_plan
+    else:
+        largest_i, largest_part = max(
+            enumerate(partitions),
+            key=lambda i_p: i_p[1].size)
+        create_part_plan.offset = largest_part.offset
+        return MultiStepPlan(plans=[
+            SlidePlan(
+                parts=[largest_part],
+                offset_delta=size),
+            ResizePlan(
+                part=largest_part,
+                size_delta=-size),
+            create_part_plan,
+            ])
+
+
+def get_boot_device_plan(device):
+    bl = device._m.bootloader
+    if bl == Bootloader.BIOS:
+        return get_boot_device_plan_bios(device)
+    if bl == Bootloader.UEFI:
+        return get_boot_device_plan_uefi(device)
+    raise Exception(f'unexpected bootloader {bl} here')
+
+
 @functools.singledispatch
 def can_be_boot_device(device, *, with_reformatting=False):
     """Can `device` be made into a boot device?
@@ -163,13 +221,13 @@ def _can_be_boot_device_disk(disk, *, with_reformatting=False):
     bl = disk._m.bootloader
     if with_reformatting:
         return True
-    if bl == Bootloader.BIOS:
+    if bl in [Bootloader.BIOS, Bootloader.UEFI]:
-        return get_boot_device_plan_bios(disk) is not None
+        return get_boot_device_plan(disk) is not None
     if disk._has_preexisting_partition():
-        if bl == Bootloader.UEFI:
+        if bl == Bootloader.PREP:
-            return any(is_esp(p) for p in disk._partitions)
-        elif bl == Bootloader.PREP:
             return any(p.flag == "prep" for p in disk._partitions)
+        else:
+            raise Exception(f'unexpected bootloader {bl} here')
     else:
         return True
 
@@ -181,10 +239,9 @@ def _can_be_boot_device_raid(raid, *, with_reformatting=False):
         return False
     if not raid.container or raid.container.metadata != 'imsm':
         return False
-    if raid._has_preexisting_partition() and not with_reformatting:
+    if with_reformatting:
-        return any(is_esp(p) for p in raid._partitions)
-    else:
         return True
+    return get_boot_device_plan_uefi(raid) is not None
 
 
 @functools.singledispatch

subiquity/common/filesystem/manipulator.py

@@ -97,14 +97,8 @@ class FilesystemManipulator:
 
     def _create_boot_partition(self, disk):
         bootloader = self.model.bootloader
-        if bootloader == Bootloader.UEFI:
+        if bootloader in [Bootloader.UEFI, Bootloader.BIOS]:
-            part_size = sizes.get_efi_size(disk)
+            boot.get_boot_device_plan(disk).apply(self)
-            log.debug('_create_boot_partition - adding EFI partition')
-            spec = dict(size=part_size, fstype='fat32')
-            if self.model._mount_for_path("/boot/efi") is None:
-                spec['mount'] = '/boot/efi'
-            self._create_boot_with_resize(
-                disk, spec, flag="boot", grub_device=True)
         elif bootloader == Bootloader.PREP:
             log.debug('_create_boot_partition - adding PReP partition')
             self._create_boot_with_resize(
@@ -113,9 +107,8 @@ class FilesystemManipulator:
                 # must be wiped or grub-install will fail
                 wipe='zero',
                 flag='prep', grub_device=True)
-        elif bootloader == Bootloader.BIOS:
+        else:
-            log.debug('_create_boot_partition - adding bios_grub partition')
+            raise Exception(f'unexpected bootloader {bootloader} here')
-            boot.get_boot_device_plan_bios(disk).apply(self)
 
     def create_raid(self, spec):
         for d in spec['devices'] | spec['spare_devices']:
@@ -340,23 +333,17 @@ class FilesystemManipulator:
         if not self.supports_resilient_boot:
             for disk in boot.all_boot_devices(self.model):
                 self.remove_boot_disk(disk)
-        if bootloader == Bootloader.BIOS:
+        if bootloader in [Bootloader.BIOS, Bootloader.UEFI]:
-            boot.get_boot_device_plan_bios(new_boot_disk).apply(self)
+            boot.get_boot_device_plan(new_boot_disk).apply(self)
             return
         if new_boot_disk._has_preexisting_partition():
-            if bootloader == Bootloader.UEFI:
+            if bootloader == Bootloader.PREP:
-                should_mount = self.model._mount_for_path('/boot/efi') is None
-                for p in new_boot_disk.partitions():
-                    if boot.is_esp(p):
-                        p.grub_device = True
-                        if should_mount:
-                            self._mount_esp(p)
-                            should_mount = False
-            elif bootloader == Bootloader.PREP:
                 for p in new_boot_disk.partitions():
                     if p.flag == 'prep':
                         p.wipe = 'zero'
                         p.grub_device = True
+            else:
+                raise Exception(f'unexpected bootloader {bootloader} here')
         else:
             if new_boot_disk.type == "disk":
                 new_boot_disk.preserve = False

subiquity/common/filesystem/tests/test_manipulator.py

@@ -18,7 +18,7 @@ import unittest
 from subiquity.common.filesystem.actions import (
     DeviceAction,
     )
-from subiquity.common.filesystem import gaps
+from subiquity.common.filesystem import boot, gaps
 from subiquity.common.filesystem.manipulator import FilesystemManipulator
 from subiquity.models.tests.test_filesystem import (
     make_disk,
@@ -285,3 +285,61 @@ class TestFilesystemManipulator(unittest.TestCase):
         self.assertIs(p2, part)
         self.assertEqual(p1.offset, MiB)
         self.assertEqual(p2.offset, half_size)
+
+    def test_add_boot_UEFI_empty(self):
+        manipulator = make_manipulator(Bootloader.UEFI)
+        disk = make_disk(manipulator.model, preserve=True)
+        manipulator.add_boot_disk(disk)
+        [part] = disk.partitions()
+        self.assertEqual(part.offset, MiB)
+
+    def test_add_boot_UEFI_full(self):
+        manipulator = make_manipulator(Bootloader.UEFI)
+        disk = make_disk(manipulator.model, preserve=True)
+        part = make_partition(
+            manipulator.model, disk, size=gaps.largest_gap_size(disk))
+        size_before = part.size
+        manipulator.add_boot_disk(disk)
+        [p1, p2] = disk.partitions()
+        self.assertIs(p2, part)
+        size_after = p2.size
+        self.assertEqual(p1.offset, MiB)
+        self.assertEqual(p2.offset, MiB+p1.size)
+        self.assertEqual(size_after, size_before - p1.size)
+
+    def test_add_boot_UEFI_half_full(self):
+        manipulator = make_manipulator(Bootloader.UEFI)
+        disk = make_disk(manipulator.model, preserve=True)
+        part = make_partition(
+            manipulator.model, disk, size=gaps.largest_gap_size(disk)//2)
+        size_before = part.size
+        manipulator.add_boot_disk(disk)
+        [p1, p2] = sorted(disk.partitions(), key=lambda p: p.offset)
+        size_after = p1.size
+        self.assertIs(p1, part)
+        self.assertEqual(p1.offset, MiB)
+        self.assertEqual(p2.offset, p1.offset + p1.size)
+        self.assertTrue(boot.is_esp(p2))
+        self.assertEqual(size_after, size_before)
+
+    def test_add_boot_UEFI_full_resizes_larger(self):
+        manipulator = make_manipulator(Bootloader.UEFI)
+        # 402MiB so that the space available for partitioning (400MiB)
+        # divided by 4 is an whole number of megabytes.
+        disk = make_disk(manipulator.model, preserve=True, size=402*MiB)
+        part_smaller = make_partition(
+            manipulator.model, disk, size=gaps.largest_gap_size(disk)//4)
+        part_larger = make_partition(
+            manipulator.model, disk, size=gaps.largest_gap_size(disk))
+        larger_size_before = part_larger.size
+        smaller_size_before = part_smaller.size
+        manipulator.add_boot_disk(disk)
+        [p1, p2, p3] = sorted(disk.partitions(), key=lambda p: p.offset)
+        self.assertIs(p1, part_smaller)
+        self.assertIs(p3, part_larger)
+        self.assertEqual(smaller_size_before, p1.size)
+        self.assertEqual(p1.offset, MiB)
+        self.assertEqual(p2.offset, p1.offset + p1.size)
+        self.assertEqual(p3.offset, p2.offset + p2.size)
+        self.assertTrue(boot.is_esp(p2))
+        self.assertEqual(p3.size, larger_size_before - p2.size)

subiquity/tests/api/test_api.py

@@ -348,27 +348,9 @@ class TestAdd(TestAPI):
             # these manual steps are expected to be equivalent to just adding
             # the single partition and getting the automatic boot partition
             await inst.post('/storage/v2/add_boot_partition', disk_id=disk_id)
-            # manual_add =
+            manual_add = await inst.post('/storage/v2/add_partition', data)
-            await inst.post('/storage/v2/add_partition', data)
 
-            # FIXME single_add places /boot/efi physically second
+            self.assertEqual(single_add, manual_add)
-            # self.assertEqual(single_add, manual_add)
-
-            # manual behaves as expected
-            # /boot/efi
-            #   offset: 1 MiB
-            #   size: 538 MiB
-            # /
-            #   offset: 539 MiB
-            #   size: 9700 MiB
-
-            # single flips the order of the physical layout
-            # /boot/efi
-            #   offset: 9701 MiB
-            #   size: 538 MiB
-            # /
-            #   offset: 1 MiB
-            #   size: 9700 MiB
 
     @timeout()
     async def test_v2_deny_multiple_add_boot_partition(self):