From f338cfed6449a31b74ae583554f7b4e679cb1104 Mon Sep 17 00:00:00 2001
From: Michael Hudson-Doyle <michael.hudson@canonical.com>
Date: Thu, 17 Aug 2023 15:21:49 +1200
Subject: [PATCH] add "in_extended" awareness to gap_with_size

---
 subiquity/common/filesystem/gaps.py           |  5 ++-
 .../common/filesystem/tests/test_gaps.py      | 38 +++++++++++++++----
 2 files changed, 33 insertions(+), 10 deletions(-)

diff --git a/subiquity/common/filesystem/gaps.py b/subiquity/common/filesystem/gaps.py
index c27f80ec..b925db8a 100644
--- a/subiquity/common/filesystem/gaps.py
+++ b/subiquity/common/filesystem/gaps.py
@@ -252,10 +252,11 @@ def largest_gap_size(device, in_extended=None):
     return 0
 
 
-def gap_with_size(device, size):
+def gap_with_size(device, size, *, in_extended=None):
     for pg in parts_and_gaps(device):
         if isinstance(pg, Gap) and pg.size >= size and pg.is_usable:
-            return pg
+            if in_extended is None or in_extended == pg.in_extended:
+                return pg
     return None
 
 
diff --git a/subiquity/common/filesystem/tests/test_gaps.py b/subiquity/common/filesystem/tests/test_gaps.py
index 6ca9dcb4..34c9e31b 100644
--- a/subiquity/common/filesystem/tests/test_gaps.py
+++ b/subiquity/common/filesystem/tests/test_gaps.py
@@ -814,11 +814,11 @@ class TestGapWithSize(GapTestCase):
             )
         )
         # 0----10---20---30---40---50---60---70---80---90---100
-        # #####     [ p1   ]                             #####
-        d = make_disk(size=100 * MiB)
-        make_partition(device=d, size=10 * MiB, offset=10 * MiB)
+        # #####     [ p1 ]                             #####
+        d = make_disk(size=100)
+        make_partition(device=d, size=10, offset=20)
         [g1, p1, g2] = gaps.parts_and_gaps(d)
-        self.assertEqual(g2, gaps.gap_with_size(d, 20 * MiB))
+        self.assertEqual(g2, gaps.gap_with_size(d, 20))
 
     def test_unusable(self):
         self.use_alignment_data(
@@ -831,7 +831,29 @@ class TestGapWithSize(GapTestCase):
             )
         )
         # 0----10---20---30---40---50---60---70---80---90---100
-        # #####     [ p1   ]                             #####
-        d = make_disk(size=100 * MiB)
-        make_partition(device=d, size=10 * MiB, offset=10 * MiB)
-        self.assertIs(None, gaps.gap_with_size(d, 10 * MiB))
+        # #####     [ p1 ]                             #####
+        d = make_disk(size=100)
+        make_partition(device=d, size=10, offset=2)
+        self.assertIs(None, gaps.gap_with_size(d, 10))
+
+    def test_in_extended(self):
+        self.use_alignment_data(
+            PartitionAlignmentData(
+                part_align=10,
+                min_gap_size=1,
+                min_start_offset=10,
+                min_end_offset=10,
+                primary_part_limit=10,
+                ebr_space=2,
+            )
+        )
+        # 0----10---20---30---40---50---60---70---80---90---100
+        # ##### g1  [ p1 (extended)          ] g3      #####
+        #            [ p5     ] g2
+        d = make_disk(size=100)
+        make_partition(device=d, size=50, offset=20, flag="extended")
+        make_partition(device=d, size=18, offset=22, flag="logical")
+        [g1, p1, p5, g2, g3] = gaps.parts_and_gaps(d)
+        self.assertEqual(g2, gaps.gap_with_size(d, 20))
+        self.assertEqual(g3, gaps.gap_with_size(d, 20, in_extended=False))
+        self.assertEqual(g2, gaps.gap_with_size(d, 10, in_extended=True))