Merge pull request #2118 from ogayot/logical-renumber

Fix logical partitions renumbering

subiquity/models/filesystem.py

@@ -731,6 +731,15 @@ def _has_preexisting_partition(self):
         return any(p.preserve for p in self._partitions)
 
     def renumber_logical_partitions(self, removed_partition):
+        """After removing a logical partition, one can call this function to
+        update the partition number of other logical partitions.
+        Please only call this function after removing a logical partition."""
+        if not removed_partition.is_logical:
+            # The implementation of this function expects the removed partition
+            # to have a partition number in the logical range (e.g., 5+ for a
+            # DOS partition table).
+            raise ValueError("do not renumber logical parts after removing a primary")
+
         parts = [
             p
             for p in self.partitions_by_number()
@@ -2233,7 +2242,8 @@ def remove_partition(self, part):
         for p2 in movable_trailing_partitions_and_gap_size(part)[0]:
             p2.offset -= part.size
         self._remove(part)
-        part.device.renumber_logical_partitions(part)
+        if part.is_logical:
+            part.device.renumber_logical_partitions(part)
         if len(part.device._partitions) == 0:
             part.device.ptable = None
 

subiquity/models/tests/test_filesystem.py

@@ -1478,6 +1478,46 @@ def test_should_add_swapfile_has_swappart(self):
         self.assertFalse(m.should_add_swapfile())
 
 
+class TestDisk(unittest.TestCase):
+    def test_renumber_logical_partitions(self):
+        m = make_model(storage_version=2)
+        d = make_disk(m, ptable="msdos")
+
+        pe = make_partition(m, d, flag="extended")
+        pl1 = make_partition(m, d, flag="logical")
+        pl2 = make_partition(m, d, flag="logical")
+        pl3 = make_partition(m, d, flag="logical")
+        pp = make_partition(m, d)
+
+        self.assertEqual(1, pe.number)
+        self.assertEqual(5, pl1.number)
+        self.assertEqual(6, pl2.number)
+        self.assertEqual(7, pl3.number)
+        self.assertEqual(2, pp.number)
+
+        d._partitions.remove(pl1)
+
+        d.renumber_logical_partitions(removed_partition=pl1)
+
+        self.assertEqual(1, pe.number)
+        self.assertEqual(5, pl2.number)
+        self.assertEqual(6, pl3.number)
+        self.assertEqual(2, pp.number)
+
+    def test_renumber_logical_partitions__after_removing_primary(self):
+        m = make_model(storage_version=2)
+        d = make_disk(m, ptable="msdos")
+
+        make_partition(m, d, flag="extended")
+        make_partition(m, d, flag="logical")
+        pp = make_partition(m, d)
+
+        d._partitions.remove(pp)
+
+        with self.assertRaisesRegex(ValueError, r"^do not renumber"):
+            d.renumber_logical_partitions(removed_partition=pp)
+
+
 class TestPartition(unittest.TestCase):
     def test_is_logical(self):
         m = make_model(storage_version=2)