ubuntu-pro: make run_answers interact with the GUI

run_answers now interacts with the GUI for the Ubuntu Pro screens.
When a contract token is provided, it is expected to be valid,
otherwise, the UI will block on Expired token or Invalid token.

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

subiquity/client/controllers/ubuntu_pro.py

@@ -18,6 +18,8 @@ import asyncio
 import logging
 from typing import Callable, List, Optional
 
+from urwid import Widget
+
 from subiquitycore.async_helpers import schedule_task
 
 from subiquity.client.controller import SubiquityTuiController
@@ -26,7 +28,12 @@ from subiquity.common.types import (
     UbuntuProCheckTokenStatus as TokenStatus,
     UbuntuProService,
     )
-from subiquity.ui.views.ubuntu_pro import UbuntuProView
+from subiquity.ui.views.ubuntu_pro import (
+    UbuntuProView,
+    UpgradeYesNoForm,
+    UpgradeModeForm,
+    TokenAddedWidget,
+    )
 
 from subiquitycore.lsb_release import lsb_release
 from subiquitycore.tuicontroller import Skip
@@ -58,9 +65,66 @@ class UbuntuProController(SubiquityTuiController):
         ubuntu_pro_info = await self.endpoint.GET()
         return UbuntuProView(self, ubuntu_pro_info.token)
 
-    def run_answers(self) -> None:
+    async def run_answers(self) -> None:
-        if "token" in self.answers:
+        """ Interact with the UI to go through the pre-attach process if
-            self.done(self.answers["token"])
+        requested. """
+        if "token" not in self.answers:
+            return
+
+        from subiquitycore.testing.view_helpers import (
+            click,
+            enter_data,
+            find_button_matching,
+            find_with_pred,
+            keypress,
+        )
+
+        view = self.app.ui.body
+
+        def run_yes_no_screen(skip: bool) -> None:
+            if skip:
+                radio = view.upgrade_yes_no_form.skip.widget
+            else:
+                radio = view.upgrade_yes_no_form.upgrade.widget
+
+            keypress(radio, key="enter")
+            click(find_button_matching(view, UpgradeYesNoForm.ok_label))
+
+        def run_token_screen(token: str) -> None:
+            keypress(view.upgrade_mode_form.with_contract_token.widget,
+                     key="enter")
+            data = {"with_contract_token_subform": {"token": token}}
+            # TODO: add this point, it would be good to trigger the validation
+            # code for the token field.
+            enter_data(view.upgrade_mode_form, data)
+            click(find_button_matching(view, UpgradeModeForm.ok_label))
+
+        async def run_token_added_overlay() -> None:
+            def is_token_added_overlay(widget: Widget) -> bool:
+                try:
+                    if widget._text == f" {TokenAddedWidget.title} ":
+                        return True
+                except AttributeError:
+                    return False
+
+            # Wait until the "Token added successfully" overlay is shown.
+            while not find_with_pred(view, is_token_added_overlay):
+                await asyncio.sleep(.2)
+
+            click(find_button_matching(view, TokenAddedWidget.done_label))
+
+        def run_services_screen() -> None:
+            click(find_button_matching(view._w,
+                                       UbuntuProView.services_done_label))
+
+        if not self.answers["token"]:
+            run_yes_no_screen(skip=True)
+            return
+
+        run_yes_no_screen(skip=False)
+        run_token_screen(self.answers["token"])
+        await run_token_added_overlay()
+        run_services_screen()
 
     def check_token(self, token: str,
                     on_success: Callable[[List[UbuntuProService]], None],

subiquity/ui/views/ubuntu_pro.py

@@ -200,6 +200,7 @@ class UbuntuProView(BaseView):
     """ Represent the view of the Ubuntu Pro configuration. """
 
     title = _("Upgrade to Ubuntu Pro")
+    services_done_label = _("Continue")
 
     def __init__(self, controller, token: str):
         """ Initialize the view with the default value for the token. """
@@ -361,7 +362,7 @@ class UbuntuProView(BaseView):
                 label=_("Back"),
                 on_press=lambda unused: on_back())
         continue_button = done_btn(
-                label=_("Continue"),
+                label=self.__class__.services_done_label,
                 on_press=lambda unused: on_continue())
 
         widgets: List[Widget] = [
@@ -538,12 +539,15 @@ class TokenAddedWidget(Stretchy):
     |                       [ Continue ]                      |
     +---------------------------------------------------------+
     """
+    title = _("Token added successfully")
+    done_label = _("Continue")
+
     def __init__(self, parent: UbuntuProView,
                  on_continue: Callable[[], None]) -> None:
         """ Initializes the widget. """
         self.parent = parent
         cont = done_btn(
-                label=_("Continue"),
+                label=self.__class__.done_label,
                 on_press=lambda unused: on_continue())
         widgets = [
             Text(_("Your token has been added successfully and your"
@@ -552,7 +556,7 @@ class TokenAddedWidget(Stretchy):
             Text(""),
             button_pile([cont]),
             ]
-        super().__init__("Token added successfully", widgets,
+        super().__init__(self.__class__.title, widgets,
                          stretchy_index=0, focus_index=2)