From ccc074d34bf34d582d8f1d43190c9024f32d9619 Mon Sep 17 00:00:00 2001
From: Olivier Gayot <olivier.gayot@canonical.com>
Date: Tue, 1 Oct 2024 14:14:26 +0200
Subject: [PATCH] storage: fix partial reformat not removing all candidate
 partitions

When reformatting a disk that has in-use (i.e., mounted) partitions, we
do a "partial" reformat, which means we attempt to delete all the
partitions that are *not* in use.

Unfortunately, the implementation was not correct. While iterating over
the partitions, calling delete_partition affects the iterator. As a
consequence, some partitions that should have been removed end up
preserved.

Fixed by iterator over a copy of the container.

LP: #2083322

Signed-off-by: Olivier Gayot <olivier.gayot@canonical.com>
---
 subiquity/server/controllers/filesystem.py    |  2 +-
 .../controllers/tests/test_filesystem.py      | 61 +++++++++++++++++++
 2 files changed, 62 insertions(+), 1 deletion(-)

diff --git a/subiquity/server/controllers/filesystem.py b/subiquity/server/controllers/filesystem.py
index da556040c..e5aba6c2e 100644
--- a/subiquity/server/controllers/filesystem.py
+++ b/subiquity/server/controllers/filesystem.py
@@ -695,7 +695,7 @@ def start_guided_reformat(
         """Perform the reformat, and return the resulting gap."""
         in_use_parts = [p for p in disk.partitions() if p._is_in_use]
         if in_use_parts:
-            for p in disk.partitions():
+            for p in list(disk.partitions()):
                 if not p._is_in_use:
                     self.delete_partition(p)
         else:
diff --git a/subiquity/server/controllers/tests/test_filesystem.py b/subiquity/server/controllers/tests/test_filesystem.py
index 13acc63ff..901294944 100644
--- a/subiquity/server/controllers/tests/test_filesystem.py
+++ b/subiquity/server/controllers/tests/test_filesystem.py
@@ -604,6 +604,67 @@ async def test__get_system_api_error_logged(self):
 
             self.assertIn("cannot load assertions for label", logs.output[0])
 
+    def test_start_guided_reformat__no_in_use(self):
+        self.fsc.model = model = make_model(Bootloader.UEFI)
+        disk = make_disk(model)
+
+        p1 = make_partition(model, disk, size=10 << 30)
+        p2 = make_partition(model, disk, size=10 << 30)
+        p3 = make_partition(model, disk, size=10 << 30)
+        p4 = make_partition(model, disk, size=10 << 30)
+
+        p_del_part = mock.patch.object(
+            self.fsc, "delete_partition", wraps=self.fsc.delete_partition
+        )
+        p_reformat = mock.patch.object(self.fsc, "reformat", wraps=self.fsc.reformat)
+
+        with p_del_part as m_del_part, p_reformat as m_reformat:
+            self.fsc.start_guided_reformat(
+                GuidedStorageTargetReformat(disk_id=disk.id), disk
+            )
+
+        m_reformat.assert_called_once_with(disk, wipe="superblock-recursive")
+        expected_del_calls = [
+            mock.call(p1, True),
+            mock.call(p2, True),
+            mock.call(p3, True),
+            mock.call(p4, True),
+        ]
+        self.assertEqual(expected_del_calls, m_del_part.mock_calls)
+
+    def test_start_guided_reformat__with_in_use(self):
+        """In LP: #2083322, start_guided_reformat did not remove all the
+        partitions that should have been removed. Because we were iterating
+        over device._partitions and calling delete_partition in the body, we
+        failed to iterate over some of the partitions."""
+        self.fsc.model = model = make_model(Bootloader.UEFI)
+        disk = make_disk(model)
+
+        p1 = make_partition(model, disk, size=10 << 30)
+        p2 = make_partition(model, disk, size=10 << 30)
+        p3 = make_partition(model, disk, size=10 << 30)
+        p4 = make_partition(model, disk, size=10 << 30)
+
+        p2._is_in_use = True
+
+        # We use wraps to ensure that the real delete_partition gets called. If
+        # we just do a no-op, we won't invalidate the iterator.
+        p_del_part = mock.patch.object(
+            self.fsc, "delete_partition", wraps=self.fsc.delete_partition
+        )
+        p_reformat = mock.patch.object(self.fsc, "reformat", wraps=self.fsc.reformat)
+
+        with p_del_part as m_del_part, p_reformat as m_reformat:
+            self.fsc.start_guided_reformat(
+                GuidedStorageTargetReformat(disk_id=disk.id), disk
+            )
+
+        m_reformat.assert_not_called()
+        # Not sure why we don't call with "override_preserve=True", like we do
+        # in reformat.
+        expected_del_calls = [mock.call(p1), mock.call(p3), mock.call(p4)]
+        self.assertEqual(expected_del_calls, m_del_part.mock_calls)
+
 
 class TestRunAutoinstallGuided(IsolatedAsyncioTestCase):
     def setUp(self):