tests: refactor shamir dry run to use input_flow

trezor · Jul 26, 2019 · 87d0174 · 87d0174
1 parent 57192a1
commit 87d0174
Showing 1 changed file with 62 additions and 65 deletions.
diff --git a/python/trezorlib/tests/device_tests/test_msg_recoverydevice_shamir_dryrun.py b/python/trezorlib/tests/device_tests/test_msg_recoverydevice_shamir_dryrun.py
@@ -2,80 +2,77 @@
 
 import pytest
 
-from trezorlib import debuglink, messages as proto
+from trezorlib import debuglink, device, messages as proto
 
-from .common import TrezorTest
+pytestmark = pytest.mark.skip_t1
 
+SHARES_20_2of3 = [
+    "crush merchant academic acid dream decision orbit smug trend trust painting slice glad crunch veteran lunch friar satoshi engage aquatic",
+    "crush merchant academic agency devote eyebrow disaster island deploy flip toxic budget numerous airport loyalty fitness resident learn sympathy daughter",
+    "crush merchant academic always course verdict rescue paces fridge museum energy solution space ladybug junction national biology game fawn coal",
+]
 
-@pytest.mark.skip_t1
-# @pytest.mark.skip(reason="waiting for shamir load_device support")
-class TestMsgRecoveryDeviceShamirDryRun(TrezorTest):
-    def test_2of3_dryrun(self):
-        # TODO: load device with these (or any other valid) mnemonics
-        mnemonics = [
-            "crush merchant academic acid dream decision orbit smug trend trust painting slice glad crunch veteran lunch friar satoshi engage aquatic",
-            "crush merchant academic agency devote eyebrow disaster island deploy flip toxic budget numerous airport loyalty fitness resident learn sympathy daughter",
-            "crush merchant academic always course verdict rescue paces fridge museum energy solution space ladybug junction national biology game fawn coal",
-        ]
-        debuglink.load_device_by_mnemonic(
-            self.client,
-            mnemonic=mnemonics[0:2],
-            pin="",
-            passphrase_protection=True,
-            label="test",
-            language="english",
-            skip_checksum=True,
-        )
 
-        ret = self.client.call_raw(
-            proto.RecoveryDevice(
-                passphrase_protection=False,
-                pin_protection=False,
-                label="label",
-                language="english",
-                dry_run=True,
-                type=proto.ResetDeviceBackupType.Slip39_Single_Group,
-            )
-        )
+def test_2of3_dryrun(client):
+    debug = client.debug
 
-        # Confirm Dryrun
-        assert isinstance(ret, proto.ButtonRequest)
-        self.client.debug.press_yes()
-        ret = self.client.call_raw(proto.ButtonAck())
+    debuglink.load_device_by_mnemonic(
+        client,
+        mnemonic=SHARES_20_2of3[0:2],
+        pin="",
+        passphrase_protection=True,
+        label="test",
+        language="english",
+        skip_checksum=True,
+    )
 
-        # Homescreen
-        assert isinstance(ret, proto.ButtonRequest)
-        self.client.debug.press_yes()
-        ret = self.client.call_raw(proto.ButtonAck())
+    def input_flow():
+        yield  # Confirm Dryrun
+        debug.press_yes()
+        # run recovery flow
+        yield from enter_all_shares(debug, SHARES_20_2of3[1:3])
 
-        # Enter word count
-        assert ret == proto.ButtonRequest(
-            code=proto.ButtonRequestType.MnemonicWordCount
+    with client:
+        client.set_input_flow(input_flow)
+        ret = device.recover(
+            client,
+            passphrase_protection=False,
+            pin_protection=False,
+            label="label",
+            language="english",
+            dry_run=True,
+            type=proto.ResetDeviceBackupType.Slip39_Single_Group,
         )
-        self.client.debug.input(str(20))
-        ret = self.client.call_raw(proto.ButtonAck())
 
-        # Homescreen
-        assert isinstance(ret, proto.ButtonRequest)
-        self.client.debug.press_yes()
-        ret = self.client.call_raw(proto.ButtonAck())
+    # Dry run was successful
+    assert ret == proto.Success(
+        message="The seed is valid and matches the one in the device"
+    )
+
 
-        # Check 2 of 3 shares
-        # TODO: check all shares when #276 is implemented
-        for mnemonic in mnemonics[1:3]:
-            assert ret == proto.ButtonRequest(
-                code=proto.ButtonRequestType.MnemonicInput
-            )
-            self.client.transport.write(proto.ButtonAck())
-            for word in mnemonic.split(" "):
-                time.sleep(1)
-                self.client.debug.input(word)
-            ret = self.client.transport.read()
+def enter_all_shares(debug, shares):
+    word_count = len(shares[0].split(" "))
 
-            # Confirm success
-            assert isinstance(ret, proto.ButtonRequest)
-            self.client.debug.press_yes()
-            ret = self.client.call_raw(proto.ButtonAck())
+    # Homescreen - proceed to word number selection
+    yield
+    debug.press_yes()
+    # Input word number
+    code = yield
+    assert code == proto.ButtonRequestType.MnemonicWordCount
+    debug.input(str(word_count))
+    # Homescreen - proceed to share entry
+    yield
+    debug.press_yes()
+    # Enter shares
+    for share in shares:
+        code = yield
+        assert code == proto.ButtonRequestType.MnemonicInput
+        # Enter mnemonic words
+        for word in share.split(" "):
+            time.sleep(1)
+            debug.input(word)
 
-        # Workflow succesfully ended
-        assert isinstance(ret, proto.Success)
+        # Homescreen - continue
+        # or Homescreen - confirm success
+        yield
+        debug.press_yes()