diff --git a/python/src/trezorlib/debuglink.py b/python/src/trezorlib/debuglink.py
index bb4fa30deb..6703184846 100644
--- a/python/src/trezorlib/debuglink.py
+++ b/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."""
diff --git a/tests/device_tests/test_busy_state.py b/tests/device_tests/test_busy_state.py
index cdd748c145..e1a0145a15 100644
--- a/tests/device_tests/test_busy_state.py
+++ b/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)