device_tests: add input flow function to Ontology tests (#348)

trezorlib/tests/device_tests/test_msg_ontology_sign_ont_id_add_attributes.py

@@ -18,7 +18,8 @@ import time
 
 import pytest
 
-from trezorlib import messages
+from trezorlib import messages, ontology
+from trezorlib.messages import ButtonRequestType as B
 from trezorlib.tools import parse_path
 
 from .common import TrezorTest
@@ -70,23 +71,28 @@ class TestMsgOntologySignOntIdAddAttributes(TrezorTest):
     def _ontology_sign(
         self, num_of_swipes, address_n, transaction, ont_id_add_attributes
     ):
-        # Sending Ontology message
-        msg = messages.OntologySignOntIdAddAttributes(
-            address_n=address_n,
-            transaction=transaction,
-            ont_id_add_attributes=ont_id_add_attributes,
-        )
-
-        self.client.transport.write(msg)
-        ret = self.client.transport.read()
+        def input_flow():
+            # Sign Tx
+            btn_code = yield
+            assert btn_code == B.SignTx
 
-        # Confirm action
-        assert isinstance(ret, messages.ButtonRequest)
-        self.client.debug.press_yes()
-        self.client.transport.write(messages.ButtonAck())
-        time.sleep(1)
-        for _ in range(num_of_swipes):
-            self.client.debug.swipe_down()
+            # Swipe and confirm
             time.sleep(1)
-        self.client.debug.press_yes()
-        return self.client.transport.read()
+            for _ in range(num_of_swipes):
+                self.client.debug.swipe_down()
+                time.sleep(1)
+
+            # Confirm Action
+            self.client.debug.press_yes()
+
+        with self.client:
+            self.client.set_input_flow(input_flow)
+            self.client.set_expected_responses(
+                [
+                    messages.ButtonRequest(code=B.SignTx),
+                    messages.OntologySignedOntIdAddAttributes(),
+                ]
+            )
+            return ontology.sign_add_attr(
+                self.client, address_n, transaction, ont_id_add_attributes
+            )

trezorlib/tests/device_tests/test_msg_ontology_sign_ont_id_register.py

@@ -18,7 +18,8 @@ import time
 
 import pytest
 
-from trezorlib import messages
+from trezorlib import messages, ontology
+from trezorlib.messages import ButtonRequestType as B
 from trezorlib.tools import parse_path
 
 from .common import TrezorTest
@@ -62,23 +63,28 @@ class TestMsgOntologySignOntIdRegister(TrezorTest):
         )
 
     def _ontology_sign(self, num_of_swipes, address_n, transaction, ont_id_register):
-        # Sending Ontology message
-        msg = messages.OntologySignOntIdRegister(
-            address_n=address_n,
-            transaction=transaction,
-            ont_id_register=ont_id_register,
-        )
-
-        self.client.transport.write(msg)
-        ret = self.client.transport.read()
+        def input_flow():
+            # Sign Tx
+            btn_code = yield
+            assert btn_code == B.SignTx
 
-        # Confirm action
-        assert isinstance(ret, messages.ButtonRequest)
-        self.client.debug.press_yes()
-        self.client.transport.write(messages.ButtonAck())
-        time.sleep(1)
-        for _ in range(num_of_swipes):
-            self.client.debug.swipe_down()
+            # Swipe and confirm
             time.sleep(1)
-        self.client.debug.press_yes()
-        return self.client.transport.read()
+            for _ in range(num_of_swipes):
+                self.client.debug.swipe_down()
+                time.sleep(1)
+
+            # Confirm Action
+            self.client.debug.press_yes()
+
+        with self.client:
+            self.client.set_expected_responses(
+                [
+                    messages.ButtonRequest(code=B.SignTx),
+                    messages.OntologySignedOntIdRegister(),
+                ]
+            )
+            self.client.set_input_flow(input_flow)
+            return ontology.sign_register(
+                self.client, address_n, transaction, ont_id_register
+            )