Merge pull request #1266 from ogayot/FR-2237

Fix storage screen sometimes not refreshing after slow probing

subiquity/client/client.py

@@ -21,7 +21,7 @@ import os
 import signal
 import sys
 import traceback
-from typing import Dict, List, Optional
+from typing import Callable, Dict, List, Optional, Union
 
 import aiohttp
 
@@ -486,11 +486,20 @@ class SubiquityClient(TuiApplication):
         log.debug("showing InstallConfirmation over %s", self.ui.body)
         self.add_global_overlay(InstallConfirmation(self))
 
-    async def _start_answers_for_view(self, controller, view):
+    async def _start_answers_for_view(
+            self, controller, view: Union[BaseView, Callable[[], BaseView]]):
+        def noop():
+            return view
+
+        if callable(view):
+            deref_view = view
+        else:
+            deref_view = noop
+
         # The view returned by make_view_for_controller is not always shown
         # immediately (if progress is being shown, but has not yet been shown
         # for a full second) so wait until it is before starting the answers.
-        while self.ui.body is not view:
+        while self.ui.body is not deref_view():
             await asyncio.sleep(0.1)
         coro = controller.run_answers()
         if inspect.iscoroutine(coro):

subiquity/client/controllers/filesystem.py

@@ -15,8 +15,10 @@
 
 import asyncio
 import logging
+from typing import Callable, Optional
 
 from subiquitycore.lsb_release import lsb_release
+from subiquitycore.view import BaseView
 
 from subiquity.client.controller import SubiquityTuiController
 from subiquity.common.filesystem import gaps
@@ -50,19 +52,34 @@ class FilesystemController(SubiquityTuiController, FilesystemManipulator):
         self.answers.setdefault('guided', False)
         self.answers.setdefault('guided-index', 0)
         self.answers.setdefault('manual', [])
+        self.current_view: Optional[BaseView] = None
+
+    async def make_ui(self) -> Callable[[], BaseView]:
+        def get_current_view() -> BaseView:
+            assert self.current_view is not None
+            return self.current_view
 
-    async def make_ui(self):
         status = await self.endpoint.guided.GET()
         if status.status == ProbeStatus.PROBING:
             self.app.aio_loop.create_task(self._wait_for_probing())
-            return SlowProbing(self)
+            self.current_view = SlowProbing(self)
         else:
-            return self.make_guided_ui(status)
+            self.current_view = self.make_guided_ui(status)
+        # NOTE: If we return a BaseView instance directly here, we have no
+        # guarantee that it will be displayed on the screen by the time the
+        # probing operation finishes. Therefore, to allow us to reliably
+        # replace the screen by the "Guided Storage" when the probing operation
+        # finishes, we add a level of indirection.
+        # In essence, this allows us to make modifications to the screen
+        # that eventually will be displayed.
+        # This is mostly a workaround for the issue described in LP #1968161
+        return get_current_view
 
     async def _wait_for_probing(self):
         status = await self.endpoint.guided.GET(wait=True)
+        self.current_view = self.make_guided_ui(status)
         if isinstance(self.ui.body, SlowProbing):
-            self.ui.set_body(self.make_guided_ui(status))
+            self.ui.set_body(self.current_view)
         else:
             log.debug("not refreshing the display. Current display is %r",
                       self.ui.body)

subiquitycore/tui.py

@@ -18,6 +18,7 @@ import inspect
 import logging
 import os
 import signal
+from typing import Callable, Optional, Union
 
 import urwid
 
@@ -34,6 +35,7 @@ from subiquitycore.tuicontroller import Skip
 from subiquitycore.ui.utils import LoadingDialog
 from subiquitycore.ui.frame import SubiquityCoreUI
 from subiquitycore.utils import astart_command
+from subiquitycore.view import BaseView
 
 log = logging.getLogger('subiquitycore.tui')
 
@@ -122,7 +124,8 @@ class TuiApplication(Application):
             before_hook()
         schedule_task(_run())
 
-    async def make_view_for_controller(self, new):
+    async def make_view_for_controller(self, new) \
+            -> Union[BaseView, Callable[[], BaseView]]:
         new.context.enter("starting UI")
         if self.opts.screens and new.name not in self.opts.screens:
             raise Skip
@@ -205,7 +208,8 @@ class TuiApplication(Application):
     async def wait_with_progress(self, awaitable):
         return await self._wait_with_indication(awaitable, self.show_progress)
 
-    async def _move_screen(self, increment, coro):
+    async def _move_screen(self, increment, coro) \
+            -> Optional[Union[BaseView, Callable[[], BaseView]]]:
         if coro is not None:
             await coro
         old, self.cur_screen = self.cur_screen, None
@@ -230,13 +234,15 @@ class TuiApplication(Application):
             except Exception:
                 self.controllers.index = cur_index
                 raise
-            else:
-                return
 
     async def move_screen(self, increment, coro):
-        view = await self.wait_with_progress(
+        view_or_callable = await self.wait_with_progress(
             self._move_screen(increment, coro))
-        if view is not None:
+        if view_or_callable is not None:
+            if callable(view_or_callable):
+                view = view_or_callable()
+            else:
+                view = view_or_callable
             self.ui.set_body(view)
 
     def next_screen(self, coro=None):