tests(core): fix upgrade tests

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()

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)
 

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:

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 (