From 2f9e3613c47e5b98611c85670ef51ff92809b353 Mon Sep 17 00:00:00 2001
From: Olivier Gayot <olivier.gayot@canonical.com>
Date: Wed, 19 Jun 2024 11:05:28 +0200
Subject: [PATCH] storage: client: accept coroutines as actions

Signed-off-by: Olivier Gayot <olivier.gayot@canonical.com>
---
 subiquity/client/controllers/filesystem.py | 22 ++++++++++++++++++++--
 1 file changed, 20 insertions(+), 2 deletions(-)

diff --git a/subiquity/client/controllers/filesystem.py b/subiquity/client/controllers/filesystem.py
index 861dcbd65..53d1a9ccc 100644
--- a/subiquity/client/controllers/filesystem.py
+++ b/subiquity/client/controllers/filesystem.py
@@ -37,7 +37,7 @@
 )
 from subiquity.ui.views import FilesystemView, GuidedDiskSelectionView
 from subiquity.ui.views.filesystem.probing import ProbingFailed, SlowProbing
-from subiquitycore.async_helpers import run_bg_task
+from subiquitycore.async_helpers import connect_async_signal, run_bg_task
 from subiquitycore.lsb_release import lsb_release
 from subiquitycore.view import BaseView
 
@@ -201,6 +201,11 @@ async def _answers_action(self, action):
         )
         from subiquitycore.ui.stretchy import StretchyOverlay
 
+        form_submitted = asyncio.Event()
+
+        async def on_form_closed(*args, **kwargs):
+            form_submitted.set()
+
         log.debug("_answers_action %r", action)
         if "obj" in action:
             obj = self._action_get(action["obj"])
@@ -214,11 +219,18 @@ async def _answers_action(self, action):
             meth = getattr(
                 self.ui.body.avail_list, "_{}_{}".format(obj.type, action_name)
             )
-            meth(obj)
+            maybe_coro = meth(obj)
+            if asyncio.iscoroutine(maybe_coro):
+                await maybe_coro
+                self.ui.body.request_redraw()
+
             yield
             body = self.ui.body._w
             if not isinstance(body, StretchyOverlay):
                 return
+
+            connect_async_signal(body.stretchy, "closed", on_form_closed)
+
             if isinstance(
                 body.stretchy, (ConfirmDeleteStretchy, ConfirmReformatStretchy)
             ):
@@ -229,10 +241,13 @@ async def _answers_action(self, action):
                     body.stretchy.form, action["data"], action.get("submit", True)
                 ):
                     pass
+
+            await form_submitted.wait()
         elif action["action"] == "create-raid":
             self.ui.body.create_raid()
             yield
             body = self.ui.body._w
+            connect_async_signal(body.stretchy, "closed", on_form_closed)
             async for _ in self._enter_form_data(
                 body.stretchy.form,
                 action["data"],
@@ -240,10 +255,12 @@ async def _answers_action(self, action):
                 clean_suffix="raid",
             ):
                 pass
+            await form_submitted.wait()
         elif action["action"] == "create-vg":
             self.ui.body.create_vg()
             yield
             body = self.ui.body._w
+            connect_async_signal(body.stretchy, "closed", on_form_closed)
             async for _ in self._enter_form_data(
                 body.stretchy.form,
                 action["data"],
@@ -251,6 +268,7 @@ async def _answers_action(self, action):
                 clean_suffix="vg",
             ):
                 pass
+            await form_submitted.wait()
         elif action["action"] == "done":
             if not self.ui.body.done_btn.enabled:
                 raise Exception("answers did not provide complete fs config")