storage: add API test to cover probert run finishing during partitioning

Signed-off-by: Olivier Gayot <olivier.gayot@canonical.com>

subiquity/common/apidef.py

@@ -277,6 +277,10 @@ def GET(wait: bool = False, use_cached_result: bool = False) \
 
         def POST(config: Payload[list]): ...
 
+        class dry_run_wait_probe:
+            """This endpoint only works in dry-run mode."""
+            def POST() -> None: ...
+
         class reset:
             def POST() -> StorageResponse: ...
 

subiquity/server/controllers/filesystem.py

@@ -866,6 +866,15 @@ async def v2_edit_partition_POST(self, data: ModifyPartitionV2) \
         self.partition_disk_handler(disk, spec, partition=partition)
         return await self.v2_GET()
 
+    async def dry_run_wait_probe_POST(self) -> None:
+        if not self.app.opts.dry_run:
+            raise NotImplementedError
+
+        # This will start the probe task if not yet started.
+        self.ensure_probing()
+
+        await self._probe_task.task
+
     @with_context(name='probe_once', description='restricted={restricted}')
     async def _probe_once(self, *, context, restricted):
         if restricted:
@@ -1049,6 +1058,14 @@ def stop_listening_udev(self):
         loop = asyncio.get_running_loop()
         loop.remove_reader(self._monitor.fileno())
 
+    def ensure_probing(self):
+        try:
+            self._probe_task.start_sync()
+        except TaskAlreadyRunningError:
+            log.debug('Skipping run of Probert - probe run already active')
+        else:
+            log.debug('Triggered Probert run on udev event')
+
     def _udev_event(self):
         cp = run_command(['udevadm', 'settle', '-t', '0'])
         if cp.returncode != 0:
@@ -1065,12 +1082,7 @@ def _udev_event(self):
         while select.select([self._monitor.fileno()], [], [], 0)[0]:
             action, dev = self._monitor.receive_device()
             log.debug("_udev_event %s %s", action, dev)
-        try:
-            self._probe_task.start_sync()
-        except TaskAlreadyRunningError:
-            log.debug('Skipping run of Probert - probe run already active')
-        else:
-            log.debug('Triggered Probert run on udev event')
+        self.ensure_probing()
 
     def make_autoinstall(self):
         rendered = self.model.render()

subiquity/tests/api/test_api.py

@@ -1533,6 +1533,49 @@ async def test_guided_v2_resize_logical_middle_partition(self):
                 self.assertLessEqual(cur['end'] + (1 << 20), nxt['offset'],
                                      f'partition overlap {cur} {nxt}')
 
+    @timeout()
+    async def test_probert_result_during_partitioning(self):
+        '''If a probert run finished during manual partition, we used to
+        load the probing data, essentially discarding changes made by the user
+        so far.  This test creates a new partition, simulates the end of a
+        probert run, and then tries to edit the previously created partition.
+        The edit operation would fail in earlier versions, because the new
+        partition would be discarded.
+        '''
+        cfg = 'examples/simple.json'
+        extra = ['--storage-version', '2']
+        async with start_server(cfg, extra_args=extra) as inst:
+            names = ['locale', 'keyboard', 'source', 'network', 'proxy',
+                     'mirror']
+            await inst.post('/meta/mark_configured', endpoint_names=names)
+            resp = await inst.get('/storage/v2')
+            [d] = resp['disks']
+            [g] = d['partitions']
+            data = {
+                'disk_id': 'disk-sda',
+                'gap': g,
+                'partition': {
+                    'size': -1,
+                    'mount': '/',
+                    'format': 'ext4',
+                }
+            }
+            add_resp = await inst.post('/storage/v2/add_partition', data)
+            [sda] = add_resp['disks']
+            [root] = match(sda['partitions'], mount='/')
+
+            # Now let's make sure we get the results from a probert run to kick
+            # in.
+            await inst.post('/storage/dry_run_wait_probe')
+            data = {
+                'disk_id': 'disk-sda',
+                'partition': {
+                    'number': root['number'],
+                }
+            }
+            # We should be able to modify the created partition.
+            await inst.post('/storage/v2/edit_partition', data)
+
 
 class TestCancel(TestAPI):
     @timeout()