From 741fca01567151fa9dad2e9d25db7236169fb697 Mon Sep 17 00:00:00 2001
From: matejcik <ja@matejcik.cz>
Date: Tue, 28 Jul 2020 14:40:54 +0200
Subject: [PATCH] tests: watch_layout must be explicit (fixes #1142)

---
 python/src/trezorlib/debuglink.py                         | 8 ++++----
 tests/device_tests/test_msg_getaddress_show.py            | 1 +
 .../device_tests/test_msg_recoverydevice_bip39_dryrun.py  | 2 ++
 tests/device_tests/test_msg_recoverydevice_bip39_t2.py    | 2 ++
 tests/device_tests/test_sdcard.py                         | 3 +++
 5 files changed, 12 insertions(+), 4 deletions(-)

diff --git a/python/src/trezorlib/debuglink.py b/python/src/trezorlib/debuglink.py
index 727dd4fae..320fa300a 100644
--- a/python/src/trezorlib/debuglink.py
+++ b/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):
diff --git a/tests/device_tests/test_msg_getaddress_show.py b/tests/device_tests/test_msg_getaddress_show.py
index 8df3d4e39..c851e92b4 100644
--- a/tests/device_tests/test_msg_getaddress_show.py
+++ b/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,
diff --git a/tests/device_tests/test_msg_recoverydevice_bip39_dryrun.py b/tests/device_tests/test_msg_recoverydevice_bip39_dryrun.py
index 1300882a8..dd577ceff 100644
--- a/tests/device_tests/test_msg_recoverydevice_bip39_dryrun.py
+++ b/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)
diff --git a/tests/device_tests/test_msg_recoverydevice_bip39_t2.py b/tests/device_tests/test_msg_recoverydevice_bip39_t2.py
index 833b24c24..e9389b91c 100644
--- a/tests/device_tests/test_msg_recoverydevice_bip39_t2.py
+++ b/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"
         )
diff --git a/tests/device_tests/test_sdcard.py b/tests/device_tests/test_sdcard.py
index 89191ed14..351dd295c 100644
--- a/tests/device_tests/test_sdcard.py
+++ b/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)