Merge pull request #2090 from ogayot/target-resize-exceeded-partitions

Don't suggest target-resize scenarios if we don't have room for partitions

subiquity/server/controllers/filesystem.py

@@ -1146,8 +1146,49 @@ def get_classic_capabilities(self):
                 classic_capabilities.update(info.capability_info.allowed)
         return sorted(classic_capabilities)
 
+    def _guided_has_enough_room_for_partitions(
+        self,
+        disk,
+        *,
+        resized_partition: Optional[Partition] = None,
+        gap: Optional[gaps.Gap] = None,
+    ) -> bool:
+        """Check if we have enough room for all the primary partitions. This
+        isn't failproof but should limit the number of TargetResize/UseGap
+        scenarios that are suggested but can't be applied because we don't have
+        enough room for partitions.
+        """
+        if (resized_partition is None) == (gap is None):
+            raise ValueError("please specify either resized_partition or gap")
+
+        new_primary_parts = 0
+        if resized_partition is not None:
+            install_to_logical = resized_partition.is_logical
+        else:
+            install_to_logical = gap.in_extended
+
+        if not install_to_logical:
+            new_primary_parts += 1
+
+        boot_plan = boot.get_boot_device_plan(disk, resize_partition=resized_partition)
+        new_primary_parts += boot_plan.new_partition_count()
+        # In theory, there could be a recovery partition as well. Not sure
+        # how to account for it since we don't know yet if one will be
+        # requested.
+        return new_primary_parts <= gaps.remaining_primary_partitions(
+            disk, disk.alignment_data()
+        )
+
+    def use_gap_has_enough_room_for_partitions(self, disk, gap: gaps.Gap) -> bool:
+        return self._guided_has_enough_room_for_partitions(disk, gap=gap)
+
+    def resize_has_enough_room_for_partitions(self, disk, resized: Partition) -> bool:
+        return self._guided_has_enough_room_for_partitions(
+            disk, resized_partition=resized
+        )
+
     def available_use_gap_scenarios(
-        self, install_min
+        self, install_min: int
     ) -> list[tuple[int, GuidedStorageTargetUseGap]]:
         scenarios: list[tuple[int, GuidedStorageTargetUseGap]] = []
         for disk in self.potential_boot_disks(with_reformatting=False):
@@ -1165,21 +1206,7 @@ def available_use_gap_scenarios(
             if gap is None:
                 # Do we return a reason here?
                 continue
-
-            # Now we check if we have enough room for all the primary
-            # partitions. This isn't failproof but should limit the number of
-            # UseGaps scenarios that are suggested but can't be applied because
-            # we don't have enough room for partitions.
-            new_primary_parts = 0
-            if not gap.in_extended:
-                new_primary_parts += 1  # For the rootfs to create.
-            new_primary_parts += boot.get_boot_device_plan(disk).new_partition_count()
-            # In theory, there could be a recovery partition as well. Not sure
-            # how to account for it since we don't know yet if one will be
-            # requested.
-            if new_primary_parts > gaps.remaining_primary_partitions(
-                disk, disk.alignment_data()
-            ):
+            if not self.use_gap_has_enough_room_for_partitions(disk, gap):
                 log.error("skipping UseGap: not enough room for primary partitions")
                 continue
 
@@ -1200,6 +1227,43 @@ def available_use_gap_scenarios(
             scenarios.append((gap.size, use_gap))
         return scenarios
 
+    def available_target_resize_scenarios(
+        self, install_min: int
+    ) -> list[tuple[int, GuidedStorageTargetResize]]:
+        scenarios: list[tuple[int, GuidedStorageTargetResize]] = []
+
+        for disk in self.potential_boot_disks(check_boot=False):
+            part_align = disk.alignment_data().part_align
+            for partition in disk.partitions():
+                if partition._is_in_use:
+                    continue
+                vals = sizes.calculate_guided_resize(
+                    partition.estimated_min_size,
+                    partition.size,
+                    install_min,
+                    part_align=part_align,
+                )
+                if vals is None:
+                    # Return a reason here
+                    continue
+                if not boot.can_be_boot_device(
+                    disk, resize_partition=partition, with_reformatting=False
+                ):
+                    # Return a reason here
+                    continue
+
+                if not self.resize_has_enough_room_for_partitions(disk, partition):
+                    log.error(
+                        "skipping TargetResize: not enough room for primary partitions"
+                    )
+                    continue
+
+                resize = GuidedStorageTargetResize.from_recommendations(
+                    partition, vals, allowed=self.get_classic_capabilities()
+                )
+                scenarios.append((vals.install_max, resize))
+        return scenarios
+
     async def v2_guided_GET(self, wait: bool = False) -> GuidedStorageResponseV2:
         """Acquire a list of possible guided storage configuration scenarios.
         Results are sorted by the size of the space potentially available to
@@ -1236,30 +1300,7 @@ async def v2_guided_GET(self, wait: bool = False) -> GuidedStorageResponseV2:
             scenarios.append((disk.size, reformat))
 
         scenarios.extend(self.available_use_gap_scenarios(install_min))
-
-        for disk in self.potential_boot_disks(check_boot=False):
-            part_align = disk.alignment_data().part_align
-            for partition in disk.partitions():
-                if partition._is_in_use:
-                    continue
-                vals = sizes.calculate_guided_resize(
-                    partition.estimated_min_size,
-                    partition.size,
-                    install_min,
-                    part_align=part_align,
-                )
-                if vals is None:
-                    # Return a reason here
-                    continue
-                if not boot.can_be_boot_device(
-                    disk, resize_partition=partition, with_reformatting=False
-                ):
-                    # Return a reason here
-                    continue
-                resize = GuidedStorageTargetResize.from_recommendations(
-                    partition, vals, allowed=classic_capabilities
-                )
-                scenarios.append((vals.install_max, resize))
+        scenarios.extend(self.available_target_resize_scenarios(install_min))
 
         scenarios.sort(reverse=True, key=lambda x: x[0])
         return GuidedStorageResponseV2(

subiquity/server/controllers/tests/test_filesystem.py

@@ -1250,17 +1250,19 @@ async def test_weighted_split(self, bootloader, ptable):
             reformat,
         )
 
-        resize = possible.pop(0)
-        expected = GuidedStorageTargetResize(
-            disk_id=self.disk.id,
-            partition_number=p.number,
-            new_size=200 << 30,
-            minimum=50 << 30,
-            recommended=200 << 30,
-            maximum=230 << 30,
-            allowed=default_capabilities,
-        )
-        self.assertEqual(expected, resize)
+        if ptable != "vtoc" or bootloader == Bootloader.NONE:
+            # VTOC has primary_part_limit=3
+            resize = possible.pop(0)
+            expected = GuidedStorageTargetResize(
+                disk_id=self.disk.id,
+                partition_number=p.number,
+                new_size=200 << 30,
+                minimum=50 << 30,
+                recommended=200 << 30,
+                maximum=230 << 30,
+                allowed=default_capabilities,
+            )
+            self.assertEqual(expected, resize)
         self.assertEqual(1, len(possible))
 
     @parameterized.expand(bootloaders_and_ptables)
@@ -1586,6 +1588,86 @@ async def test_available_use_gap_scenarios(
 
         self.assertEqual(expected_scenario, scenarios != [])
 
+    async def test_resize_has_enough_room_for_partitions__one_primary(self):
+        await self._setup(Bootloader.NONE, "gpt", fix_bios=True)
+
+        p = make_partition(self.model, self.disk, preserve=True, size=4 << 20)
+
+        self.assertTrue(self.fsc.resize_has_enough_room_for_partitions(self.disk, p))
+
+    async def test_resize_has_enough_room_for_partitions__full_primaries(self):
+        await self._setup(Bootloader.NONE, "dos", fix_bios=True)
+
+        p1 = make_partition(self.model, self.disk, preserve=True, size=4 << 20)
+        p2 = make_partition(self.model, self.disk, preserve=True, size=4 << 20)
+        p3 = make_partition(self.model, self.disk, preserve=True, size=4 << 20)
+        p4 = make_partition(self.model, self.disk, preserve=True, size=4 << 20)
+
+        self.assertFalse(self.fsc.resize_has_enough_room_for_partitions(self.disk, p1))
+        self.assertFalse(self.fsc.resize_has_enough_room_for_partitions(self.disk, p2))
+        self.assertFalse(self.fsc.resize_has_enough_room_for_partitions(self.disk, p3))
+        self.assertFalse(self.fsc.resize_has_enough_room_for_partitions(self.disk, p4))
+
+    @mock.patch("subiquity.server.controllers.filesystem.boot.get_boot_device_plan")
+    async def test_resize_has_enough_room_for_partitions__one_more(self, p_boot_plan):
+        await self._setup(Bootloader.NONE, "dos", fix_bios=True)
+
+        model = self.model
+        disk = self.disk
+        # Sizes are irrelevant
+        size = 4 << 20
+        p1 = make_partition(model, disk, preserve=True, size=size)
+        p2 = make_partition(model, disk, preserve=True, size=size)
+        p3 = make_partition(model, disk, preserve=True, size=size)
+
+        p_boot_plan.return_value = boot.CreatePartPlan(
+            mock.Mock(), mock.Mock(), mock.Mock()
+        )
+        self.assertFalse(self.fsc.resize_has_enough_room_for_partitions(disk, p1))
+        self.assertFalse(self.fsc.resize_has_enough_room_for_partitions(disk, p2))
+        self.assertFalse(self.fsc.resize_has_enough_room_for_partitions(disk, p3))
+
+        p_boot_plan.return_value = boot.NoOpBootPlan()
+        self.assertTrue(self.fsc.resize_has_enough_room_for_partitions(disk, p1))
+        self.assertTrue(self.fsc.resize_has_enough_room_for_partitions(disk, p2))
+        self.assertTrue(self.fsc.resize_has_enough_room_for_partitions(disk, p3))
+
+    @mock.patch("subiquity.server.controllers.filesystem.boot.get_boot_device_plan")
+    async def test_resize_has_enough_room_for_partitions__logical(self, p_boot_plan):
+        await self._setup(Bootloader.NONE, "dos", fix_bios=True)
+
+        model = self.model
+        disk = self.disk
+        # Sizes are irrelevant
+        size = 4 << 20
+        p1 = make_partition(model, disk, preserve=True, size=size)
+        p2 = make_partition(model, disk, preserve=True, size=size, flag="extended")
+        p5 = make_partition(model, disk, preserve=True, size=size, flag="logical")
+        p6 = make_partition(model, disk, preserve=True, size=size, flag="logical")
+        p3 = make_partition(model, disk, preserve=True, size=size)
+        p4 = make_partition(model, disk, preserve=True, size=size)
+
+        p_boot_plan.return_value = boot.CreatePartPlan(
+            mock.Mock(), mock.Mock(), mock.Mock()
+        )
+        self.assertFalse(self.fsc.resize_has_enough_room_for_partitions(disk, p1))
+        self.assertFalse(self.fsc.resize_has_enough_room_for_partitions(disk, p2))
+        self.assertFalse(self.fsc.resize_has_enough_room_for_partitions(disk, p3))
+        self.assertFalse(self.fsc.resize_has_enough_room_for_partitions(disk, p4))
+        # we're installing in a logical partition, but we still have not enough
+        # room to apply the boot plan.
+        self.assertFalse(self.fsc.resize_has_enough_room_for_partitions(disk, p5))
+        self.assertFalse(self.fsc.resize_has_enough_room_for_partitions(disk, p6))
+
+        p_boot_plan.return_value = boot.NoOpBootPlan()
+        self.assertFalse(self.fsc.resize_has_enough_room_for_partitions(disk, p1))
+        self.assertFalse(self.fsc.resize_has_enough_room_for_partitions(disk, p2))
+        self.assertFalse(self.fsc.resize_has_enough_room_for_partitions(disk, p3))
+        self.assertFalse(self.fsc.resize_has_enough_room_for_partitions(disk, p4))
+        # if we're installing in a logical partition, we have enough room
+        self.assertTrue(self.fsc.resize_has_enough_room_for_partitions(disk, p5))
+        self.assertTrue(self.fsc.resize_has_enough_room_for_partitions(disk, p6))
+
 
 class TestManualBoot(IsolatedAsyncioTestCase):
     def _setup(self, bootloader, ptable, **kw):