diff --git a/subiquity/client/controllers/ubuntu_pro.py b/subiquity/client/controllers/ubuntu_pro.py index 703377dd..678ec434 100644 --- a/subiquity/client/controllers/ubuntu_pro.py +++ b/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: - if "token" in self.answers: - self.done(self.answers["token"]) + async def run_answers(self) -> None: + """ Interact with the UI to go through the pre-attach process if + 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], diff --git a/subiquity/ui/views/ubuntu_pro.py b/subiquity/ui/views/ubuntu_pro.py index 9b09020d..3d49bdad 100644 --- a/subiquity/ui/views/ubuntu_pro.py +++ b/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)