adapt make_boot_disk for existing boot partitions

subiquity/controllers/filesystem.py

@@ -484,11 +484,23 @@ class FilesystemController(BaseController):
 
     def make_boot_disk(self, new_boot_disk):
         boot_partition = None
-        for disk in self.model.all_disks():
+        if self.model.bootloader == Bootloader.BIOS:
-            for part in disk.partitions():
+            install_dev = self.model.grub_install_device
-                if part.flag in ("bios_grub", "boot", "prep"):
+            if install_dev:
-                    boot_partition = part
+                boot_partition = install_dev._potential_boot_partition()
+        elif self.model.bootloader == Bootloader.UEFI:
+            mount = self.model._mount_for_path("/boot/efi")
+            if mount is not None:
+                boot_partition = mount.device.volume
+        elif self.model.bootloader == Bootloader.PREP:
+            boot_partition = self.model.grub_install_device
         if boot_partition is not None:
+            if boot_partition.preserve:
+                if self.model.bootloader == Bootloader.PREP:
+                    boot_partition.wipe = None
+                elif self.model.bootloader == Bootloader.UEFI:
+                    self.delete_mount(boot_partition.fs().mount())
+            else:
                 boot_disk = boot_partition.device
                 full = boot_disk.free_for_partitions == 0
                 self.delete_partition(boot_partition)
@@ -500,5 +512,19 @@ class FilesystemController(BaseController):
                     largest_part = max(
                         new_boot_disk.partitions(), key=lambda p: p.size)
                     largest_part.size -= (
-                    boot_partition.size - new_boot_disk.free_for_partitions)
+                        boot_partition.size -
+                        new_boot_disk.free_for_partitions)
+        if new_boot_disk._has_preexisting_partition():
+            if self.model.bootloader == Bootloader.BIOS:
+                self.model.grub_install_device = new_boot_disk
+            elif self.model.bootloader == Bootloader.UEFI:
+                part = new_boot_disk._potential_boot_partition()
+                if part.fs() is None:
+                    self.model.add_filesystem(part, 'fat32')
+                self.model.add_mount(part.fs(), '/boot/efi')
+            elif self.model.bootloader == Bootloader.PREP:
+                part = new_boot_disk._potential_boot_partition()
+                part.wipe = 'zero'
+                self.model.grub_install_device = part
+        else:
             self._create_boot_partition(new_boot_disk)

subiquity/controllers/tests/test_filesystem.py

@@ -83,8 +83,8 @@ class TestFilesystemController(unittest.TestCase):
 
     def test_make_boot_disk_BIOS(self):
         controller = make_controller(Bootloader.BIOS)
-        disk1 = make_disk(controller.model)
+        disk1 = make_disk(controller.model, preserve=False)
-        disk2 = make_disk(controller.model)
+        disk2 = make_disk(controller.model, preserve=False)
         disk2p1 = controller.model.add_partition(
             disk2, size=disk2.free_for_partitions)
 
@@ -103,10 +103,28 @@ class TestFilesystemController(unittest.TestCase):
         self.assertEqual(disk2.partitions()[0].flag, "bios_grub")
         self.assertEqual(controller.model.grub_install_device, disk2)
 
+    def test_make_boot_disk_BIOS_existing(self):
+        controller = make_controller(Bootloader.BIOS)
+        disk1 = make_disk(controller.model, preserve=True)
+        disk1p1 = controller.model.add_partition(
+            disk1, size=1 << 20, flag="bios_grub")
+        disk1p1.preserve = True
+        disk2 = make_disk(controller.model, preserve=False)
+
+        self.assertEqual(disk1.partitions(), [disk1p1])
+        self.assertEqual(controller.model.grub_install_device, None)
+        controller.make_boot_disk(disk1)
+        self.assertEqual(disk1.partitions(), [disk1p1])
+        self.assertEqual(controller.model.grub_install_device, disk1)
+
+        controller.make_boot_disk(disk2)
+        self.assertEqual(disk1.partitions(), [disk1p1])
+        self.assertEqual(controller.model.grub_install_device, disk2)
+
     def test_make_boot_disk_UEFI(self):
         controller = make_controller(Bootloader.UEFI)
-        disk1 = make_disk(controller.model)
+        disk1 = make_disk(controller.model, preserve=False)
-        disk2 = make_disk(controller.model)
+        disk2 = make_disk(controller.model, preserve=False)
         disk2p1 = controller.model.add_partition(
             disk2, size=disk2.free_for_partitions)
 
@@ -130,10 +148,35 @@ class TestFilesystemController(unittest.TestCase):
         efi_mnt = controller.model._mount_for_path("/boot/efi")
         self.assertEqual(efi_mnt.device.volume, disk2.partitions()[0])
 
+    def test_make_boot_disk_UEFI_existing(self):
+        controller = make_controller(Bootloader.UEFI)
+        disk1 = make_disk(controller.model, preserve=True)
+        disk1p1 = controller.model.add_partition(
+            disk1, size=512 << 20, flag="boot")
+        disk1p1.preserve = True
+        disk2 = make_disk(controller.model, preserve=True)
+
+        self.assertEqual(disk1.partitions(), [disk1p1])
+        self.assertEqual(controller.model.grub_install_device, None)
+        efi_mnt = controller.model._mount_for_path("/boot/efi")
+        self.assertEqual(efi_mnt, None)
+        controller.make_boot_disk(disk1)
+        self.assertEqual(disk1.partitions(), [disk1p1])
+        self.assertEqual(controller.model.grub_install_device, None)
+        efi_mnt = controller.model._mount_for_path("/boot/efi")
+        self.assertEqual(efi_mnt.device.volume, disk1p1)
+        self.assertEqual(disk1p1.fs().fstype, "fat32")
+
+        controller.make_boot_disk(disk2)
+        self.assertEqual(disk1.partitions(), [disk1p1])
+        self.assertEqual(controller.model.grub_install_device, None)
+        efi_mnt = controller.model._mount_for_path("/boot/efi")
+        self.assertEqual(efi_mnt.device.volume, disk2.partitions()[0])
+
     def test_make_boot_disk_PREP(self):
         controller = make_controller(Bootloader.PREP)
-        disk1 = make_disk(controller.model)
+        disk1 = make_disk(controller.model, preserve=False)
-        disk2 = make_disk(controller.model)
+        disk2 = make_disk(controller.model, preserve=False)
         disk2p1 = controller.model.add_partition(
             disk2, size=disk2.free_for_partitions)
 
@@ -157,3 +200,28 @@ class TestFilesystemController(unittest.TestCase):
         self.assertEqual(
             controller.model.grub_install_device,
             disk2.partitions()[0])
+
+    def test_make_boot_disk_PREP_existing(self):
+        controller = make_controller(Bootloader.PREP)
+        disk1 = make_disk(controller.model, preserve=True)
+        disk1p1 = controller.model.add_partition(
+            disk1, size=8 << 20, flag="prep")
+        disk1p1.preserve = True
+        disk2 = make_disk(controller.model, preserve=False)
+
+        self.assertEqual(disk1.partitions(), [disk1p1])
+        self.assertEqual(controller.model.grub_install_device, None)
+        controller.make_boot_disk(disk1)
+        self.assertEqual(disk1.partitions(), [disk1p1])
+        self.assertEqual(controller.model.grub_install_device, disk1p1)
+        self.assertEqual(disk1p1.wipe, 'zero')
+
+        controller.make_boot_disk(disk2)
+        self.assertEqual(disk1.partitions(), [disk1p1])
+        self.assertEqual(disk1p1.wipe, None)
+        self.assertEqual(
+            controller.model.grub_install_device, disk2.partitions()[0])
+        self.assertEqual(disk2.partitions()[0].flag, "prep")
+        self.assertEqual(
+            controller.model.grub_install_device,
+            disk2.partitions()[0])