From 0f78185b9a7456e84dfa27c4c9d459b449100218 Mon Sep 17 00:00:00 2001
From: Michael Hudson-Doyle <michael.hudson@canonical.com>
Date: Wed, 28 Jun 2023 13:43:52 +1200
Subject: [PATCH] add some unit tests

---
 subiquity/server/controllers/filesystem.py    |  7 +++--
 .../controllers/tests/test_filesystem.py      | 31 +++++++++++++++++++
 2 files changed, 36 insertions(+), 2 deletions(-)

diff --git a/subiquity/server/controllers/filesystem.py b/subiquity/server/controllers/filesystem.py
index c84b940f..d0b5bc7c 100644
--- a/subiquity/server/controllers/filesystem.py
+++ b/subiquity/server/controllers/filesystem.py
@@ -125,6 +125,9 @@ system_non_gpt_text = _(
 )
 
 
+DRY_RUN_RESET_SIZE = 500*MiB
+
+
 class NoSnapdSystemsOnSource(Exception):
     pass
 
@@ -566,11 +569,11 @@ class FilesystemController(SubiquityController, FilesystemManipulator):
 
         if choice.reset_partition:
             if self.app.opts.dry_run:
-                reset_size = 500*MiB
+                reset_size = DRY_RUN_RESET_SIZE
             else:
                 cp = await arun_command(['du', '-sb', '/cdrom'])
                 reset_size = int(cp.stdout.strip().split()[0])
-            reset_size = align_up(int(reset_size * 1.10), 256 * MiB)
+                reset_size = align_up(int(reset_size * 1.10), 256 * MiB)
             reset_gap, gap = gap.split(reset_size)
             self.reset_partition = self.create_partition(
                 device=reset_gap.device, gap=reset_gap,
diff --git a/subiquity/server/controllers/tests/test_filesystem.py b/subiquity/server/controllers/tests/test_filesystem.py
index 724979a3..5c5927aa 100644
--- a/subiquity/server/controllers/tests/test_filesystem.py
+++ b/subiquity/server/controllers/tests/test_filesystem.py
@@ -54,6 +54,7 @@ from subiquity.models.tests.test_filesystem import (
     )
 from subiquity.server import snapdapi
 from subiquity.server.controllers.filesystem import (
+    DRY_RUN_RESET_SIZE,
     FilesystemController,
     VariationInfo,
     )
@@ -340,6 +341,36 @@ class TestGuided(IsolatedAsyncioTestCase):
         self.assertFalse(d1p2.preserve)
         self.assertIsNone(gaps.largest_gap(self.d1))
 
+    async def test_guided_reset_partition(self):
+        await self._guided_setup(Bootloader.UEFI, 'gpt')
+        target = GuidedStorageTargetReformat(
+            disk_id=self.d1.id, allowed=default_capabilities)
+        await self.controller.guided(
+            GuidedChoiceV2(
+                target=target,
+                capability=GuidedCapability.DIRECT,
+                reset_partition=True))
+        [d1p1, d1p2, d1p3] = self.d1.partitions()
+        self.assertEqual('/boot/efi', d1p1.mount)
+        self.assertEqual(None, d1p2.mount)
+        self.assertEqual(DRY_RUN_RESET_SIZE, d1p2.size)
+        self.assertEqual('/', d1p3.mount)
+
+    async def test_guided_reset_partition_only(self):
+        await self._guided_setup(Bootloader.UEFI, 'gpt')
+        target = GuidedStorageTargetReformat(
+            disk_id=self.d1.id, allowed=default_capabilities)
+        await self.controller.guided(
+            GuidedChoiceV2(
+                target=target,
+                capability=GuidedCapability.DIRECT,
+                reset_partition=True),
+            reset_partition_only=True)
+        [d1p1, d1p2] = self.d1.partitions()
+        self.assertEqual(None, d1p1.mount)
+        self.assertEqual(None, d1p2.mount)
+        self.assertEqual(DRY_RUN_RESET_SIZE, d1p2.size)
+
     async def test_guided_direct_BIOS_MSDOS(self):
         await self._guided_setup(Bootloader.BIOS, 'msdos')
         target = GuidedStorageTargetReformat(