tests: only allow input_flow in a with-block

python/src/trezorlib/debuglink.py

@@ -237,6 +237,9 @@ class TrezorClientDebugLink(TrezorClient):
             self.ui.input_flow = None
             return
 
+        if not self.in_with_statement:
+            raise RuntimeError("Must be called inside 'with' statement")
+
         if callable(input_flow):
             input_flow = input_flow()
         if not hasattr(input_flow, "send"):
@@ -252,6 +255,9 @@ class TrezorClientDebugLink(TrezorClient):
     def __exit__(self, _type, value, traceback):
         self.in_with_statement -= 1
 
+        # Clear input flow.
+        self.set_input_flow(None)
+
         if _type is not None:
             # Another exception raised
             return False
@@ -268,6 +274,7 @@ class TrezorClientDebugLink(TrezorClient):
         # Cleanup
         self.expected_responses = None
         self.current_response = None
+
         return False
 
     def set_expected_responses(self, expected):

tests/device_tests/test_msg_tezos_sign_tx.py

@@ -194,24 +194,26 @@ class TestMsgTezosSignTx:
         debug.press_yes()
 
     def test_tezos_sign_tx_proposal(self, client):
-        client.set_input_flow(self.input_flow(client.debug, num_pages=1))
-        resp = tezos.sign_tx(
-            client,
-            TEZOS_PATH_10,
-            dict_to_proto(
-                messages.TezosSignTx,
-                {
-                    "branch": "dee04042c0832d68a43699b2001c0a38065436eb05e578071a763e1972d0bc81",
-                    "proposal": {
-                        "source": "005f450441f41ee11eee78a31d1e1e55627c783bd6",
-                        "period": 17,
-                        "proposals": [
-                            "dfa974df171c2dad9a9b8f25d99af41fd9702ce5d04521d2f9943c84d88aa572"
-                        ],
+        with client:
+            client.set_input_flow(self.input_flow(client.debug, num_pages=1))
+            resp = tezos.sign_tx(
+                client,
+                TEZOS_PATH_10,
+                dict_to_proto(
+                    messages.TezosSignTx,
+                    {
+                        "branch": "dee04042c0832d68a43699b2001c0a38065436eb05e578071a763e1972d0bc81",
+                        "proposal": {
+                            "source": "005f450441f41ee11eee78a31d1e1e55627c783bd6",
+                            "period": 17,
+                            "proposals": [
+                                "dfa974df171c2dad9a9b8f25d99af41fd9702ce5d04521d2f9943c84d88aa572"
+                            ],
+                        },
                     },
-                },
-            ),
-        )
+                ),
+            )
+
         assert (
             resp.signature
             == "edsigtfY16R32k2WVMYfFr7ymnro4ib5zMckk28vsuViYNN77DJAvCJLRNArd9L531pUCxT4YdcvCvBym5dhcZ1rknEVm6yZ8bB"
@@ -225,25 +227,27 @@ class TestMsgTezosSignTx:
         )
 
     def test_tezos_sign_tx_multiple_proposals(self, client):
-        client.set_input_flow(self.input_flow(client.debug, num_pages=2))
-        resp = tezos.sign_tx(
-            client,
-            TEZOS_PATH_10,
-            dict_to_proto(
-                messages.TezosSignTx,
-                {
-                    "branch": "7e0be36a90c663c73c60da3889ffefff1383fb65cc29f0639f173d8f95a52df7",
-                    "proposal": {
-                        "source": "005f450441f41ee11eee78a31d1e1e55627c783bd6",
-                        "period": 17,
-                        "proposals": [
-                            "2a6ff28ab4d0ccb18f7129aaaf9a4b8027d794f2562849665fdb6999db2a4e57",
-                            "47cd60c09ab8437cc9fe19add494dce1b9844100f660f02ce77510a0c66d2762",
-                        ],
+        with client:
+            client.set_input_flow(self.input_flow(client.debug, num_pages=2))
+            resp = tezos.sign_tx(
+                client,
+                TEZOS_PATH_10,
+                dict_to_proto(
+                    messages.TezosSignTx,
+                    {
+                        "branch": "7e0be36a90c663c73c60da3889ffefff1383fb65cc29f0639f173d8f95a52df7",
+                        "proposal": {
+                            "source": "005f450441f41ee11eee78a31d1e1e55627c783bd6",
+                            "period": 17,
+                            "proposals": [
+                                "2a6ff28ab4d0ccb18f7129aaaf9a4b8027d794f2562849665fdb6999db2a4e57",
+                                "47cd60c09ab8437cc9fe19add494dce1b9844100f660f02ce77510a0c66d2762",
+                            ],
+                        },
                     },
-                },
-            ),
-        )
+                ),
+            )
+
         assert (
             resp.signature
             == "edsigu6GAjhiWAQ64ctWTGEDYAZ16tYzLgzWzqc4CUyixK4FGRE8YUBVzFaVJ2fUCexZjZLMLdiNZGcUdzeL1bQhZ2h5oLrh7pA"

tests/device_tests/test_reset_recovery_slip39_advanced.py

@@ -179,7 +179,6 @@ def reset(client, strength=128):
             language="english",
             backup_type=BackupType.Slip39_Advanced,
         )
-    client.set_input_flow(None)
 
     # Check if device is properly initialized
     assert client.features.initialized is True
@@ -203,8 +202,6 @@ def recover(client, shares):
         client.set_input_flow(input_flow)
         ret = device.recover(client, pin_protection=False, label="label")
 
-    client.set_input_flow(None)
-
     # Workflow successfully ended
     assert ret == messages.Success(message="Device recovered")
     assert client.features.pin_protection is False

tests/device_tests/test_reset_recovery_slip39_basic.py

@@ -115,8 +115,6 @@ def reset(client, strength=128):
             backup_type=BackupType.Slip39_Basic,
         )
 
-    client.set_input_flow(None)
-
     # Check if device is properly initialized
     assert client.features.initialized is True
     assert client.features.needs_backup is False
@@ -139,8 +137,6 @@ def recover(client, shares):
         client.set_input_flow(input_flow)
         ret = device.recover(client, pin_protection=False, label="label")
 
-    client.set_input_flow(None)
-
     # Workflow successfully ended
     assert ret == messages.Success(message="Device recovered")
     assert client.features.pin_protection is False