tests: watch_layout must be explicit (fixes #1142)

python/src/trezorlib/debuglink.py

@@ -22,6 +22,7 @@ from mnemonic import Mnemonic
 
 from . import mapping, messages, protobuf
 from .client import TrezorClient
+from .exceptions import TrezorFailure
 from .log import DUMP_BYTES
 from .tools import expect
 
@@ -85,6 +86,8 @@ class DebugLink:
 
     def wait_layout(self):
         obj = self._call(messages.DebugLinkGetState(wait_layout=True))
+        if isinstance(obj, messages.Failure):
+            raise TrezorFailure(obj)
         return layout_lines(obj.layout_lines)
 
     def watch_layout(self, watch: bool) -> None:
@@ -344,11 +347,9 @@ class TrezorClientDebugLink(TrezorClient):
         self.ui.input_flow = input_flow
         input_flow.send(None)  # start the generator
 
-    def watch_layout(self, watch: bool) -> None:
+    def watch_layout(self, watch: bool = True) -> None:
         """Enable or disable watching layout changes.
 
-        Happens implicitly in a `with client` block.
-
         Since trezor-core v2.3.2, it is necessary to call `watch_layout()` before
         using `debug.wait_layout()`, otherwise layout changes are not reported.
         """
@@ -362,7 +363,6 @@ class TrezorClientDebugLink(TrezorClient):
     def __enter__(self):
         # For usage in with/expected_responses
         self.in_with_statement += 1
-        self.watch_layout(True)
         return self
 
     def __exit__(self, _type, value, traceback):

tests/device_tests/test_msg_getaddress_show.py

@@ -141,6 +141,7 @@ def test_show_multisig_xpubs(client):
             client.debug.press_yes()
 
         with client:
+            client.watch_layout()
             client.set_input_flow(input_flow)
             btc.get_address(
                 client,

tests/device_tests/test_msg_recoverydevice_bip39_dryrun.py

@@ -81,6 +81,7 @@ def do_recover_core(client, mnemonic, **kwargs):
         client.debug.click(buttons.OK)
 
     with client:
+        client.watch_layout()
         client.set_input_flow(input_flow)
         return device.recover(client, dry_run=True, **kwargs)
 
@@ -160,6 +161,7 @@ def test_invalid_seed_core(client):
         client.debug.click(buttons.OK)
 
     with client:
+        client.watch_layout()
         client.set_input_flow(input_flow)
         with pytest.raises(exceptions.Cancelled):
             return device.recover(client, dry_run=True)

tests/device_tests/test_msg_recoverydevice_bip39_t2.py

@@ -64,6 +64,7 @@ def test_tt_pin_passphrase(client):
 
     with client:
         client.set_input_flow(input_flow)
+        client.watch_layout()
         device.recover(
             client, pin_protection=True, passphrase_protection=True, label="hello"
         )
@@ -109,6 +110,7 @@ def test_tt_nopin_nopassphrase(client):
 
     with client:
         client.set_input_flow(input_flow)
+        client.watch_layout()
         device.recover(
             client, pin_protection=False, passphrase_protection=False, label="hello"
         )

tests/device_tests/test_sdcard.py

@@ -66,6 +66,7 @@ def test_sd_protect_unlock(client):
         client.debug.press_yes()
 
     with client:
+        client.watch_layout()
         client.set_input_flow(input_flow_enable_sd_protect)
         device.sd_protect(client, Op.ENABLE)
 
@@ -91,6 +92,7 @@ def test_sd_protect_unlock(client):
         client.debug.press_yes()
 
     with client:
+        client.watch_layout()
         client.set_input_flow(input_flow_change_pin)
         device.change_pin(client)
 
@@ -110,6 +112,7 @@ def test_sd_protect_unlock(client):
         client.debug.press_no()  # close
 
     with client, pytest.raises(TrezorFailure) as e:
+        client.watch_layout()
         client.set_input_flow(input_flow_change_pin_format)
         device.change_pin(client)