make boot part planning for UEFI & PREP v2-aware

subiquity/common/filesystem/boot.py

@@ -194,12 +194,12 @@ def get_boot_device_plan_bios(device) -> Optional[MakeBootDevicePlan]:
         enumerate(movable),
         key=lambda i_p: i_p[1].size)
     return MultiStepPlan(plans=[
-        SlidePlan(
-            parts=movable[:largest_i+1],
-            offset_delta=sizes.BIOS_GRUB_SIZE_BYTES),
         ResizePlan(
             part=largest_part,
             size_delta=-sizes.BIOS_GRUB_SIZE_BYTES),
+        SlidePlan(
+            parts=movable[:largest_i+1],
+            offset_delta=sizes.BIOS_GRUB_SIZE_BYTES),
         create_part_plan,
         attr_plan,
         ])
@@ -216,32 +216,31 @@ def get_add_part_plan(device, *, spec, args):
         create_part_plan.offset = gaps.largest_gap(device).offset
         return create_part_plan
     else:
-        largest_i, largest_part = max(
+        new_parts = [p for p in partitions if not p.preserve]
-            enumerate(partitions),
+        if not new_parts:
-            key=lambda i_p: i_p[1].size)
+            return None
+        largest_part = max(new_parts, key=lambda p: p.size)
         if size > largest_part.size // 2:
             return None
         create_part_plan.offset = largest_part.offset
         return MultiStepPlan(plans=[
-            SlidePlan(
-                parts=[largest_part],
-                offset_delta=size),
             ResizePlan(
                 part=largest_part,
                 size_delta=-size),
+            SlidePlan(
+                parts=[largest_part],
+                offset_delta=size),
             create_part_plan,
             ])
 
 
 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
 
     spec = dict(size=sizes.get_efi_size(device), fstype='fat32', mount=None)
     if device._m._mount_for_path("/boot/efi") is None:
@@ -252,14 +251,12 @@ def get_boot_device_plan_uefi(device):
 
 
 def get_boot_device_plan_prep(device):
-    if device._has_preexisting_partition():
     for part in device.partitions():
         if part.flag == "prep":
             return MultiStepPlan(plans=[
                 SetAttrPlan(part, 'grub_device', True),
                 SetAttrPlan(part, 'wipe', 'zero')
                 ])
-        return None
 
     return get_add_part_plan(
         device,

subiquity/common/filesystem/tests/test_manipulator.py

@@ -582,21 +582,98 @@ class TestFilesystemManipulator(unittest.TestCase):
             manipulator.add_boot_disk(disk)
         self.assertIsBootDisk(manipulator, disk)
 
-    def test_add_boot_UEFI_full_resizes_larger(self):
+    @parameterized.expand([
-        manipulator = make_manipulator(Bootloader.UEFI)
+        (Bootloader.UEFI, 1),
-        # 2002MiB so that the space available for partitioning (2000MiB)
+        (Bootloader.PREP, 2),
-        # divided by 4 is an whole number of megabytes.
+        (Bootloader.UEFI, 1),
-        disk = make_disk(manipulator.model, preserve=True, size=2002*MiB)
+        (Bootloader.PREP, 2)])
-        self._test_add_boot_full_resizes_larger(
+    def test_no_add_boot_full_preserved_partition(self, bl, version):
-            manipulator, disk, sizes.get_efi_size(disk))
+        manipulator = make_manipulator(bl, version)
+        disk = make_disk(manipulator.model, preserve=True)
+        full_size = gaps.largest_gap_size(disk)
+        make_partition(
+            manipulator.model, disk, size=full_size, preserve=True)
+        self.assertFalse(boot.can_be_boot_device(disk))
 
-    def test_add_boot_PREP_full_resizes_larger(self):
+    @parameterized.expand([
-        manipulator = make_manipulator(Bootloader.PREP)
+        (Bootloader.UEFI, 1),
-        # 2002MiB so that the space available for partitioning (2000MiB)
+        (Bootloader.PREP, 2),
-        # divided by 4 is an whole number of megabytes.
+        (Bootloader.UEFI, 1),
+        (Bootloader.PREP, 2)])
+    def test_add_boot_half_preserved_partition(self, bl, version):
+        manipulator = make_manipulator(bl, version)
+        disk = make_disk(manipulator.model, preserve=True)
+        half_size = gaps.largest_gap_size(disk)//2
+        part = make_partition(
+            manipulator.model, disk, size=half_size,
+            preserve=True)
+        size = self.boot_size_for_disk(disk)
+        if version == 1:
+            self.assertFalse(boot.can_be_boot_device(disk))
+        else:
+            with self.assertPartitionOperations(
+                    disk,
+                    Unchanged(part),
+                    Create(
+                        offset=part.offset + part.size,
+                        size=size),
+                    ):
+                manipulator.add_boot_disk(disk)
+            self.assertIsBootDisk(manipulator, disk)
+
+    @parameterized.expand([
+        (Bootloader.UEFI,),
+        (Bootloader.PREP,)])
+    def test_add_boot_half_preserved_half_new_partition(self, bl):
+        # This test is v2 only because you can only get into this
+        # situation in a v2 world.
+        manipulator = make_manipulator(bl, 2)
+        disk = make_disk(manipulator.model, preserve=True)
+        half_size = gaps.largest_gap_size(disk)//2
+        old_part = make_partition(
+            manipulator.model, disk, size=half_size)
+        new_part = make_partition(
+            manipulator.model, disk, size=half_size)
+        old_part.preserve = True
+        size = self.boot_size_for_disk(disk)
+        with self.assertPartitionOperations(
+                disk,
+                Unchanged(old_part),
+                Create(
+                    offset=new_part.offset,
+                    size=size),
+                MoveResize(
+                    part=new_part,
+                    offset=size,
+                    size=-size),
+                ):
+            manipulator.add_boot_disk(disk)
+        self.assertIsBootDisk(manipulator, disk)
+
+    @parameterized.expand([
+        (Bootloader.UEFI, 1),
+        (Bootloader.PREP, 2),
+        (Bootloader.UEFI, 1),
+        (Bootloader.PREP, 2)])
+    def test_add_boot_existing_boot_partition(self, bl, version):
+        manipulator = make_manipulator(bl, version)
         disk = make_disk(manipulator.model, preserve=True, size=2002*MiB)
-        self._test_add_boot_full_resizes_larger(
+        if bl == Bootloader.UEFI:
-            manipulator, disk, sizes.PREP_GRUB_SIZE_BYTES)
+            flag = 'boot'
+        else:
+            flag = 'prep'
+        boot_part = make_partition(
+            manipulator.model, disk, size=self.boot_size_for_disk(disk),
+            flag=flag)
+        rest_part = make_partition(manipulator.model, disk)
+        boot_part.preserve = rest_part.preserve = True
+        with self.assertPartitionOperations(
+                disk,
+                Unchanged(boot_part),
+                Unchanged(rest_part),
+                ):
+            manipulator.add_boot_disk(disk)
+        self.assertIsBootDisk(manipulator, disk)
 
 
 class TestReformat(unittest.TestCase):