diff --git a/python/src/trezorlib/debuglink.py b/python/src/trezorlib/debuglink.py index 3219b2c1c..a6f77b6fd 100644 --- a/python/src/trezorlib/debuglink.py +++ b/python/src/trezorlib/debuglink.py @@ -460,6 +460,11 @@ class DebugLink: """Differences in handling debug events and LayoutContent.""" return self.version < (2, 6, 1) + @property + def has_global_layout(self) -> bool: + """Differences in waiting for Global Layout objects.""" + return self.version >= (2, 8, 6) + @property def layout_type(self) -> LayoutType: assert self.model is not None @@ -517,9 +522,13 @@ class DebugLink: self._write(msg) return self._read() - def state( - self, wait_type: DebugWaitType = DebugWaitType.CURRENT_LAYOUT - ) -> messages.DebugLinkState: + def state(self, wait_type: DebugWaitType | None = None) -> messages.DebugLinkState: + if wait_type is None: + wait_type = ( + DebugWaitType.CURRENT_LAYOUT + if self.has_global_layout + else DebugWaitType.IMMEDIATE + ) result = self._call(messages.DebugLinkGetState(wait_layout=wait_type)) while not isinstance(result, (messages.Failure, messages.DebugLinkState)): result = self._read() @@ -554,6 +563,9 @@ class DebugLink: ["DUMMY CONTENT, WAIT UNTIL THE END OF THE BLOCK :("] ) + # make sure some current layout is up by issuing a dummy GetState + self.state() + # send GetState without waiting for reply self._write(messages.DebugLinkGetState(wait_layout=DebugWaitType.NEXT_LAYOUT)) @@ -563,6 +575,7 @@ class DebugLink: yield layout_content finally: self.waiting_for_layout_change = False + self.layout_dirty = True # wait for the reply resp = self._read() diff --git a/tests/device_handler.py b/tests/device_handler.py index 4baf8e87e..45ec1df9f 100644 --- a/tests/device_handler.py +++ b/tests/device_handler.py @@ -37,9 +37,10 @@ class NullUI: class BackgroundDeviceHandler: _pool = ThreadPoolExecutor() - def __init__(self, client: "Client") -> None: + def __init__(self, client: "Client", nowait: bool = False) -> None: self._configure_client(client) self.task = None + self.nowait = nowait def _configure_client(self, client: "Client") -> None: self.client = client @@ -55,11 +56,7 @@ class BackgroundDeviceHandler: if self.task is not None: raise RuntimeError("Wait for previous task first") - # make sure we start the wait while a layout is up - # TODO should this be part of "wait_for_layout_change"? - self.debuglink().read_layout() - # from the displayed layout, wait for the first UI change triggered by the - # task running in the background + # wait for the first UI change triggered by the task running in the background with self.debuglink().wait_for_layout_change(): self.task = self._pool.submit(function, self.client, *args, **kwargs) diff --git a/tests/upgrade_tests/recovery_old.py b/tests/upgrade_tests/recovery_old.py index 5dedc18bf..6b8fed4c2 100644 --- a/tests/upgrade_tests/recovery_old.py +++ b/tests/upgrade_tests/recovery_old.py @@ -11,9 +11,9 @@ def _enter_word( ) -> "LayoutContent": typed_word = word[:4] for coords in buttons.type_word(typed_word, is_slip39=is_slip39): - debug.click(coords) + debug.click(coords, wait=False) - return debug.click(buttons.CONFIRM_WORD) + return debug.click(buttons.CONFIRM_WORD, wait=True) def confirm_recovery(debug: "DebugLink") -> None: diff --git a/tests/upgrade_tests/test_firmware_upgrades.py b/tests/upgrade_tests/test_firmware_upgrades.py index d74bea90c..bafe67f51 100644 --- a/tests/upgrade_tests/test_firmware_upgrades.py +++ b/tests/upgrade_tests/test_firmware_upgrades.py @@ -301,7 +301,7 @@ def test_upgrade_shamir_recovery(gen: str, tag: Optional[str]): device_handler.run(device.recover, pin_protection=False) recovery_old.confirm_recovery(debug) - recovery_old.select_number_of_words(debug, wait=not debug.legacy_debug) + recovery_old.select_number_of_words(debug) layout = recovery_old.enter_share(debug, MNEMONIC_SLIP39_BASIC_20_3of6[0]) if not debug.legacy_ui and not debug.legacy_debug: assert (