From cb2d184ef2edf2ae114598289d4ce85368fead3b Mon Sep 17 00:00:00 2001
From: Dan Bungert <daniel.bungert@canonical.com>
Date: Fri, 10 Jun 2022 15:50:56 -0600
Subject: [PATCH] gaps: add gap.split

---
 subiquity/common/filesystem/gaps.py           | 21 +++++++++++++++++-
 .../common/filesystem/tests/test_gaps.py      | 22 +++++++++++++++++++
 2 files changed, 42 insertions(+), 1 deletion(-)

diff --git a/subiquity/common/filesystem/gaps.py b/subiquity/common/filesystem/gaps.py
index d212ae65..d4581724 100644
--- a/subiquity/common/filesystem/gaps.py
+++ b/subiquity/common/filesystem/gaps.py
@@ -28,7 +28,8 @@ from subiquity.models.filesystem import (
     )
 
 
-@attr.s(auto_attribs=True)
+# should also set on_setattr=None with attrs 20.1.0
+@attr.s(auto_attribs=True, frozen=True)
 class Gap:
     device: object
     offset: int
@@ -41,6 +42,24 @@ class Gap:
     def id(self):
         return 'gap-' + self.device.id
 
+    def split(self, size):
+        """returns a tuple of two new gaps, split from the current gap based on
+        the supplied size.  If size is equal to the gap size, the second gap is
+        None.  The original gap is unmodified."""
+        if size > self.size:
+            raise Exception('requested size larger than gap')
+        if size == self.size:
+            return (self, None)
+        first_gap = Gap(device=self.device,
+                        offset=self.offset,
+                        size=size,
+                        in_extended=self.in_extended)
+        rest_gap = Gap(device=self.device,
+                       offset=self.offset + size,
+                       size=self.size - size,
+                       in_extended=self.in_extended)
+        return (first_gap, rest_gap)
+
 
 @functools.singledispatch
 def parts_and_gaps(device):
diff --git a/subiquity/common/filesystem/tests/test_gaps.py b/subiquity/common/filesystem/tests/test_gaps.py
index 04fc5344..c346b43e 100644
--- a/subiquity/common/filesystem/tests/test_gaps.py
+++ b/subiquity/common/filesystem/tests/test_gaps.py
@@ -37,6 +37,28 @@ class TestGaps(unittest.TestCase):
         self.assertEqual(MiB, gap.offset)
 
 
+class TestSplitGap(unittest.TestCase):
+    def test_equal(self):
+        [gap] = gaps.parts_and_gaps(make_disk())
+        actual = gap.split(gap.size)
+        self.assertEqual((gap, None), actual)
+
+    def test_too_big(self):
+        [gap] = gaps.parts_and_gaps(make_disk())
+        with self.assertRaises(Exception):
+            gap.split(gap.size + MiB)
+
+    def test_split(self):
+        [gap] = gaps.parts_and_gaps(make_disk(size=100 << 30))
+        size = 10 << 30
+        new_gaps = gap.split(size)
+        self.assertEqual(2, len(new_gaps))
+        self.assertEqual(size, new_gaps[0].size)
+        self.assertEqual(gap.size - size, new_gaps[1].size)
+        self.assertEqual(gap.offset, new_gaps[0].offset)
+        self.assertEqual(gap.offset + size, new_gaps[1].offset)
+
+
 class TestDiskGaps(unittest.TestCase):
 
     def test_no_partition_gpt(self):