From 3789b3196012bc63ad3a35a99c06ecdf6ba6f1ec Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Yao=20Wei=20=28=E9=AD=8F=E9=8A=98=E5=BB=B7=29?=
 <yao.wei@canonical.com>
Date: Mon, 18 Dec 2023 11:17:23 +0100
Subject: [PATCH] reset_partition accepts fixed size

---
 subiquity/common/types.py                     |  1 +
 subiquity/server/controllers/filesystem.py    | 23 ++++++++++-
 .../controllers/tests/test_filesystem.py      | 40 +++++++++++++++++++
 3 files changed, 62 insertions(+), 2 deletions(-)

diff --git a/subiquity/common/types.py b/subiquity/common/types.py
index 712f2699c2..08d909b9b3 100644
--- a/subiquity/common/types.py
+++ b/subiquity/common/types.py
@@ -531,6 +531,7 @@ class GuidedChoiceV2:
 
     sizing_policy: Optional[SizingPolicy] = SizingPolicy.SCALED
     reset_partition: bool = False
+    reset_partition_size: Optional[int] = None
 
 
 @attr.s(auto_attribs=True)
diff --git a/subiquity/server/controllers/filesystem.py b/subiquity/server/controllers/filesystem.py
index 9995e4e990..cd16ff16c0 100644
--- a/subiquity/server/controllers/filesystem.py
+++ b/subiquity/server/controllers/filesystem.py
@@ -29,6 +29,7 @@
 from curtin import swap
 from curtin.commands.extract import AbstractSourceHandler
 from curtin.storage_config import ptable_part_type_to_flag
+from curtin.util import human2bytes
 
 from subiquity.common.apidef import API
 from subiquity.common.errorreport import ErrorReportKind
@@ -707,7 +708,10 @@ async def guided(
             raise Exception("failed to locate gap after adding boot")
 
         if choice.reset_partition:
-            if self.app.opts.dry_run:
+            if choice.reset_partition_size is not None:
+                part_align = disk.alignment_data().part_align
+                reset_size = align_up(choice.reset_partition_size, part_align)
+            elif self.app.opts.dry_run:
                 reset_size = DRY_RUN_RESET_SIZE
             else:
                 cp = await arun_command(["du", "-sb", "/cdrom"])
@@ -1405,6 +1409,20 @@ async def run_autoinstall_guided(self, layout):
                 disk_id=gap.device.id, gap=gap, allowed=[]
             )
 
+        reset_partition = False
+        reset_partition_size = None
+        rp_input = layout.get("reset-partition", None)
+        if rp_input:
+            reset_partition = True
+            if type(rp_input) in (str, int):
+                reset_partition_size = int(human2bytes(rp_input))
+                log.info(
+                    "autoinstall: will install reset partition"
+                    f"of size {reset_partition_size}"
+                )
+            else:
+                log.info("autoinstall: will install reset partition")
+
         log.info(
             f"autoinstall: running guided {capability} install in "
             f"mode {mode} using {target}"
@@ -1416,7 +1434,8 @@ async def run_autoinstall_guided(self, layout):
                 password=password,
                 recovery_key=guided_recovery_key,
                 sizing_policy=sizing_policy,
-                reset_partition=layout.get("reset-partition", False),
+                reset_partition=reset_partition,
+                reset_partition_size=reset_partition_size,
             ),
             reset_partition_only=layout.get("reset-partition-only", False),
         )
diff --git a/subiquity/server/controllers/tests/test_filesystem.py b/subiquity/server/controllers/tests/test_filesystem.py
index 4321601df1..d75b83d9e7 100644
--- a/subiquity/server/controllers/tests/test_filesystem.py
+++ b/subiquity/server/controllers/tests/test_filesystem.py
@@ -459,6 +459,46 @@ async def test_guided_reset_partition(self):
         self.assertEqual(DRY_RUN_RESET_SIZE, d1p2.size)
         self.assertEqual("/", d1p3.mount)
 
+    async def test_fixed_reset_partition(self):
+        await self._guided_setup(Bootloader.UEFI, "gpt")
+        target = GuidedStorageTargetReformat(
+            disk_id=self.d1.id, allowed=default_capabilities
+        )
+        fixed_reset_size = 12 << 30
+        await self.controller.guided(
+            GuidedChoiceV2(
+                target=target,
+                capability=GuidedCapability.DIRECT,
+                reset_partition=True,
+                reset_partition_size=fixed_reset_size,
+            )
+        )
+        [d1p1, d1p2, d1p3] = self.d1.partitions()
+        self.assertEqual("/boot/efi", d1p1.mount)
+        self.assertEqual(None, d1p2.mount)
+        self.assertEqual(fixed_reset_size, d1p2.size)
+        self.assertEqual("/", d1p3.mount)
+
+    async def test_fixed_reset_partition_human(self):
+        await self._guided_setup(Bootloader.UEFI, "gpt")
+        target = GuidedStorageTargetReformat(
+            disk_id=self.d1.id, allowed=default_capabilities
+        )
+        fixed_reset_size = "12G"
+        await self.controller.guided(
+            GuidedChoiceV2(
+                target=target,
+                capability=GuidedCapability.DIRECT,
+                reset_partition=True,
+                reset_partition_size=fixed_reset_size,
+            )
+        )
+        [d1p1, d1p2, d1p3] = self.d1.partitions()
+        self.assertEqual("/boot/efi", d1p1.mount)
+        self.assertEqual(None, d1p2.mount)
+        self.assertEqual(12 << 30, d1p2.size)
+        self.assertEqual("/", d1p3.mount)
+
     async def test_guided_reset_partition_only(self):
         await self._guided_setup(Bootloader.UEFI, "gpt")
         target = GuidedStorageTargetReformat(