From 2abdfb5a4880ad7d1f5156668187903670f26b54 Mon Sep 17 00:00:00 2001
From: Michael Hudson-Doyle <michael.hudson@canonical.com>
Date: Thu, 17 Aug 2023 14:31:31 +1200
Subject: [PATCH] add gap_with_size function to find a gap of a given size

---
 subiquity/common/filesystem/gaps.py           |  7 +++
 .../common/filesystem/tests/test_gaps.py      | 52 +++++++++++++++++++
 subiquity/models/tests/test_filesystem.py     |  5 +-
 3 files changed, 63 insertions(+), 1 deletion(-)

diff --git a/subiquity/common/filesystem/gaps.py b/subiquity/common/filesystem/gaps.py
index 736b4240..c27f80ec 100644
--- a/subiquity/common/filesystem/gaps.py
+++ b/subiquity/common/filesystem/gaps.py
@@ -252,6 +252,13 @@ def largest_gap_size(device, in_extended=None):
     return 0
 
 
+def gap_with_size(device, size):
+    for pg in parts_and_gaps(device):
+        if isinstance(pg, Gap) and pg.size >= size and pg.is_usable:
+            return pg
+    return None
+
+
 @functools.singledispatch
 def movable_trailing_partitions_and_gap_size(partition):
     """For a given partition (or LVM logical volume), return the total,
diff --git a/subiquity/common/filesystem/tests/test_gaps.py b/subiquity/common/filesystem/tests/test_gaps.py
index 6277dfd8..6ca9dcb4 100644
--- a/subiquity/common/filesystem/tests/test_gaps.py
+++ b/subiquity/common/filesystem/tests/test_gaps.py
@@ -783,3 +783,55 @@ class TestUsable(unittest.TestCase):
         self.assertEqual(
             "TOO_MANY_PRIMARY_PARTS", GapUsable.TOO_MANY_PRIMARY_PARTS.name
         )
+
+
+class TestGapWithSize(GapTestCase):
+    def test_empty_disk(self):
+        d = make_disk(size=10 * MiB)
+        [g1] = gaps.parts_and_gaps(d)
+        self.assertEqual(g1, gaps.gap_with_size(d, MiB))
+
+    def test_half_full(self):
+        d = make_disk(size=10 * MiB)
+        make_partition(device=d, size=d.size // 2)
+        [p1, g1] = gaps.parts_and_gaps(d)
+        self.assertEqual(g1, gaps.gap_with_size(d, MiB))
+
+    def test_half_full_too_big(self):
+        d = make_disk(size=10 * MiB)
+        make_partition(device=d, size=d.size // 2)
+        [p1, g1] = gaps.parts_and_gaps(d)
+        self.assertIs(None, gaps.gap_with_size(d, 10 * MiB))
+
+    def test_one_gap_too_small(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,
+            )
+        )
+        # 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)
+        [g1, p1, g2] = gaps.parts_and_gaps(d)
+        self.assertEqual(g2, gaps.gap_with_size(d, 20 * MiB))
+
+    def test_unusable(self):
+        self.use_alignment_data(
+            PartitionAlignmentData(
+                part_align=10,
+                min_gap_size=1,
+                min_start_offset=10,
+                min_end_offset=10,
+                primary_part_limit=1,
+            )
+        )
+        # 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))
diff --git a/subiquity/models/tests/test_filesystem.py b/subiquity/models/tests/test_filesystem.py
index e0762cec..4688f615 100644
--- a/subiquity/models/tests/test_filesystem.py
+++ b/subiquity/models/tests/test_filesystem.py
@@ -164,10 +164,13 @@ def make_model_and_disk(bootloader=None, **kw):
     return model, make_disk(model, **kw)
 
 
-def make_partition(model, device=None, *, preserve=False, size=None, offset=None, **kw):
+def make_partition(
+    model=None, device=None, *, preserve=False, size=None, offset=None, **kw
+):
     flag = kw.pop("flag", None)
     if device is None:
         device = make_disk(model)
+    model = device._m
     if size is None or offset is None:
         gap = gaps.largest_gap(device)
         if size is None: