diff --git a/subiquity/common/filesystem/gaps.py b/subiquity/common/filesystem/gaps.py index f569e71a..5cbfb4a2 100644 --- a/subiquity/common/filesystem/gaps.py +++ b/subiquity/common/filesystem/gaps.py @@ -241,3 +241,12 @@ def within(device, gap): if pg.offset >= gap.offset and pg_end <= gap_end: return pg return None + + +def after(device, offset): + """Find the first gap that is after this offset.""" + for pg in parts_and_gaps(device): + if isinstance(pg, Gap): + if pg.offset > offset: + return pg + return None diff --git a/subiquity/common/filesystem/tests/test_gaps.py b/subiquity/common/filesystem/tests/test_gaps.py index 08c5677e..4d8e2749 100644 --- a/subiquity/common/filesystem/tests/test_gaps.py +++ b/subiquity/common/filesystem/tests/test_gaps.py @@ -135,6 +135,35 @@ class TestWithin(unittest.TestCase): self.assertEqual(g2, gaps.within(d, orig_g2)) +class TestAfter(unittest.TestCase): + def test_basic(self): + d = make_disk() + [gap] = gaps.parts_and_gaps(d) + self.assertEqual(gap, gaps.after(d, 0)) + + def test_no_equals(self): + d = make_disk() + [gap] = gaps.parts_and_gaps(d) + self.assertIsNone(gaps.after(d, gap.offset)) + + def test_partition_resize_full_part(self): + m, d = make_model_and_disk() + [g1] = gaps.parts_and_gaps(d) + p1 = make_partition(m, d, size=g1.size) + p1.size //= 2 + gap = gaps.after(d, p1.offset) + self.assertIsNotNone(gap) + + def test_partition_resize_half_part(self): + m, d = make_model_and_disk() + make_partition(m, d) + [p1, g1] = gaps.parts_and_gaps(d) + p1.size //= 2 + gap = gaps.after(d, p1.offset) + self.assertNotEqual(gap, g1) + self.assertTrue(gap.offset < g1.offset) + + class TestDiskGaps(unittest.TestCase): def test_no_partition_gpt(self):