filesystem: only use logical part renum if forced

subiquity/models/filesystem.py

@@ -519,6 +519,9 @@ class _Device(_Formattable, ABC):
     def partitions_by_offset(self):
         return sorted(self._partitions, key=lambda p: p.offset)
 
+    def partitions_by_number(self):
+        return sorted(self._partitions, key=lambda p: p.number)
+
     @property
     def used(self):
         if self._is_entirely_used():
@@ -671,10 +674,10 @@ class Disk(_Device):
             return None
         return id.encode('utf-8').decode('unicode_escape').strip()
 
-    def renumber_logical_partitions(self):
+    def renumber_logical_partitions(self, removed_partition):
-        parts = [p for p in self.partitions_by_offset() if p.is_logical]
+        parts = [p for p in self.partitions_by_number()
-        primary_limit = self.alignment_data().primary_part_limit
+                 if p.is_logical and p.number > removed_partition.number]
-        next_num = primary_limit + 1
+        next_num = removed_partition.number
         for part in parts:
             part.number = next_num
             next_num += 1
@@ -700,17 +703,19 @@ class Partition(_Formattable):
         if self.number is not None:
             return
 
-        if self.is_logical:
-            self.device.renumber_logical_partitions()
-        else:
         used_nums = {p.number for p in self.device._partitions
-                         if p.number is not None and not p.is_logical}
+                     if p.number is not None
+                     if p.is_logical == self.is_logical}
         primary_limit = self.device.alignment_data().primary_part_limit
-            for num in range(1, primary_limit + 1):
+        if self.is_logical:
+            possible_nums = range(primary_limit + 1, 129)
+        else:
+            possible_nums = range(1, primary_limit + 1)
+        for num in possible_nums:
             if num not in used_nums:
                 self.number = num
                 return
-            raise Exception('Exceeded number of available primary partitions')
+        raise Exception('Exceeded number of available partitions')
 
     def available(self):
         if self.flag in ['bios_grub', 'prep'] or self.grub_device:
@@ -1466,7 +1471,7 @@ class FilesystemModel(object):
         for p2 in movable_trailing_partitions_and_gap_size(part)[0]:
             p2.offset -= part.size
         self._remove(part)
-        part.device.renumber_logical_partitions()
+        part.device.renumber_logical_partitions(part)
         if len(part.device._partitions) == 0:
             part.device.ptable = None
 

subiquity/models/tests/test_filesystem.py

@@ -739,6 +739,32 @@ class TestPartitionNumbering(unittest.TestCase):
         for p in parts:
             self.assert_next(p)
 
+    @parameterized.expand(
+        [[pt, primaries, i]
+         for pt, primaries in (('msdos', 4), ('vtoc', 3))
+         for i in range(3)]
+    )
+    def test_out_of_offset_order(self, ptable, primaries, idx_to_remove):
+        m = make_model(storage_version=2)
+        d1 = make_disk(m, ptable=ptable, size=100 << 20)
+        self.assert_next(make_partition(m, d1, size=10 << 20))
+        self.assert_next(make_partition(m, d1, flag='extended'))
+        self.cur_idx = primaries + 1
+        parts = []
+        parts.append(make_partition(
+                m, d1, flag='logical', size=9 << 20, offset=30 << 20))
+        parts.append(make_partition(
+                m, d1, flag='logical', size=9 << 20, offset=20 << 20))
+        parts.append(make_partition(
+                m, d1, flag='logical', size=9 << 20, offset=40 << 20))
+        for p in parts:
+            self.assert_next(p)
+        to_remove = parts.pop(idx_to_remove)
+        m.remove_partition(to_remove)
+        self.cur_idx = primaries + 1
+        for p in parts:
+            self.assert_next(p)
+
     @parameterized.expand([
         [1, 'msdos', 4],
         [1, 'vtoc', 3],