fix(tests): fix unstable UI in test_busy_expiry testcase

Asserting the device comes back to Homescreen, so the UI result is always the same
(when it does not get there, the assert will fail and the test will rerun).

[no changelog]

python/src/trezorlib/debuglink.py

@@ -77,7 +77,7 @@ class UnstructuredJSONReader:
             self.dict = {}
 
     def top_level_value(self, key: str) -> Any:
-        return self.dict[key]
+        return self.dict.get(key)
 
     def find_objects_with_key_and_value(self, key: str, value: Any) -> List["AnyDict"]:
         def recursively_find(data: Any) -> Iterator[Any]:
@@ -140,7 +140,7 @@ class LayoutContent(UnstructuredJSONReader):
 
     def main_component(self) -> str:
         """Getting the main component of the layout."""
-        return self.top_level_value("component")
+        return self.top_level_value("component") or "no main component"
 
     def all_components(self) -> List[str]:
         """Getting all components of the layout."""
@@ -332,7 +332,7 @@ class LayoutContent(UnstructuredJSONReader):
 
     def tt_pin_digits_order(self) -> str:
         """In what order the PIN buttons are shown on the screen. Only for TT."""
-        return self.top_level_value("digits_order")
+        return self.top_level_value("digits_order") or "no digits order"
 
     def get_middle_choice(self) -> str:
         """What is the choice being selected right now."""

tests/device_tests/test_busy_state.py

@@ -25,15 +25,24 @@ from trezorlib.tools import parse_path
 PIN = "1234"
 
 
+def _assert_busy(client: Client, should_be_busy: bool, screen: str = "Homescreen"):
+    assert client.features.busy is should_be_busy
+    if client.debug.model in ("T", "R"):
+        if should_be_busy:
+            assert "CoinJoinProgress" in client.debug.read_layout().all_components()
+        else:
+            assert client.debug.read_layout().main_component() == screen
+
+
 @pytest.mark.setup_client(pin=PIN)
 def test_busy_state(client: Client):
+    _assert_busy(client, False, "Lockscreen")
     assert client.features.unlocked is False
-    assert client.features.busy is False
 
     # Show busy dialog for 1 minute.
     device.set_busy(client, expiry_ms=60 * 1000)
+    _assert_busy(client, True)
     assert client.features.unlocked is False
-    assert client.features.busy is True
 
     with client:
         client.use_pin_sequence([PIN])
@@ -42,24 +51,31 @@ def test_busy_state(client: Client):
         )
 
     client.refresh_features()
+    _assert_busy(client, True)
     assert client.features.unlocked is True
-    assert client.features.busy is True
 
     # Hide the busy dialog.
     device.set_busy(client, None)
 
+    _assert_busy(client, False)
     assert client.features.unlocked is True
-    assert client.features.busy is False
 
 
 @pytest.mark.flaky(max_runs=5)
 def test_busy_expiry(client: Client):
+    _assert_busy(client, False)
     # Show the busy dialog.
-    device.set_busy(client, expiry_ms=100)
+    device.set_busy(client, expiry_ms=1500)
+    _assert_busy(client, True)
 
-    # Wait for it to expire. Add 100ms tolerance to account for CI slowness.
-    time.sleep(0.2)
+    # Hasn't expired yet.
+    time.sleep(1.4)
+    _assert_busy(client, True)
+
+    # Wait for it to expire. Add 400ms tolerance to account for CI slowness.
+    time.sleep(0.5)
 
     # Check that the device is no longer busy.
+    # Also needs to come back to Homescreen (for UI tests).
     client.refresh_features()
-    assert client.features.busy is False
+    _assert_busy(client, False)