tests/persistence: extend recovery test to test a full recovery

including restarts

tests/click_tests/recovery.py

@@ -8,11 +8,14 @@ def enter_word(debug, word):
     return debug.click(buttons.CONFIRM_WORD, wait=True)
 
 
-def select_number_of_words(debug, num_of_words=20):
+def confirm_recovery(debug):
-    # confirm recovery
     layout = debug.wait_layout()
     assert layout.text.startswith("Recovery mode")
-    layout = debug.click(buttons.OK, wait=True)
+    debug.click(buttons.OK, wait=True)
+
+
+def select_number_of_words(debug, num_of_words=20):
+    layout = debug.read_layout()
 
     # select number of words
     assert "Select number of words" in layout.text

tests/click_tests/test_recovery.py

@@ -31,6 +31,8 @@ def test_recovery(device_handler):
     assert features.initialized is False
     device_handler.run(device.recover, pin_protection=False)
 
+    recovery.confirm_recovery(debug)
+
     recovery.select_number_of_words(debug)
     recovery.enter_shares(debug, MNEMONIC_SLIP39_BASIC_20_3of6)
     recovery.finalize(debug)

tests/persistence_tests/test_shamir_persistence.py

@@ -19,18 +19,13 @@ import pytest
 from trezorlib import device
 
 from .. import buttons
+from ..click_tests import recovery
+from ..common import MNEMONIC_SLIP39_ADVANCED_20, MNEMONIC_SLIP39_BASIC_20_3of6
 from ..device_handler import BackgroundDeviceHandler
 from ..emulators import EmulatorWrapper
 from ..upgrade_tests import core_only
 
 
-def enter_word(debug, word):
-    word = word[:4]
-    for coords in buttons.type_word(word):
-        debug.click(coords)
-    return debug.click(buttons.CONFIRM_WORD, wait=True)
-
-
 @pytest.fixture
 def emulator():
     emu = EmulatorWrapper("core")
@@ -38,8 +33,13 @@ def emulator():
         yield emu
 
 
+def _restart(device_handler: BackgroundDeviceHandler, emulator: EmulatorWrapper):
+    device_handler.restart(emulator)
+    return device_handler.debuglink()
+
+
 @core_only
-def test_persistence(emulator):
+def test_abort(emulator):
     device_handler = BackgroundDeviceHandler(emulator.client)
     debug = device_handler.debuglink()
     features = device_handler.features()
@@ -70,3 +70,80 @@ def test_persistence(emulator):
     assert layout.text == "Homescreen"
     features = device_handler.features()
     assert features.recovery_mode is False
+
+
+@core_only
+def test_recovery_single_reset(emulator):
+    device_handler = BackgroundDeviceHandler(emulator.client)
+    debug = device_handler.debuglink()
+    features = device_handler.features()
+
+    assert features.initialized is False
+    assert features.recovery_mode is False
+
+    device_handler.run(device.recover, pin_protection=False)
+    recovery.confirm_recovery(debug)
+
+    recovery.select_number_of_words(debug)
+
+    debug = _restart(device_handler, emulator)
+    features = device_handler.features()
+    assert features.recovery_mode is True
+
+    # we need to enter the number of words again, that's a feature
+    recovery.select_number_of_words(debug)
+    recovery.enter_shares(debug, MNEMONIC_SLIP39_BASIC_20_3of6)
+    recovery.finalize(debug)
+
+    features = device_handler.features()
+    assert features.initialized is True
+    assert features.recovery_mode is False
+
+
+@core_only
+def test_recovery_multiple_resets(emulator):
+    def enter_shares_with_restarts(debug):
+        shares = MNEMONIC_SLIP39_ADVANCED_20
+        layout = debug.read_layout()
+        expected_text = "Enter any share"
+        remaining = len(shares)
+        for share in shares:
+            assert expected_text in layout.text
+            layout = recovery.enter_share(debug, share)
+            remaining -= 1
+            expected_text = "Success You have entered"
+            debug = _restart(device_handler, emulator)
+
+        assert "You have successfully recovered your wallet" in layout.text
+
+    device_handler = BackgroundDeviceHandler(emulator.client)
+    debug = device_handler.debuglink()
+    features = device_handler.features()
+
+    assert features.initialized is False
+    assert features.recovery_mode is False
+
+    # start device and recovery
+    device_handler.run(device.recover, pin_protection=False)
+    recovery.confirm_recovery(debug)
+
+    # set number of words
+    recovery.select_number_of_words(debug)
+
+    # restart
+    debug = _restart(device_handler, emulator)
+    features = device_handler.features()
+    assert features.recovery_mode is True
+
+    # enter the number of words again, that's a feature!
+    recovery.select_number_of_words(debug)
+
+    # enter shares and restart after each one
+    enter_shares_with_restarts(debug)
+    debug = device_handler.debuglink()
+    layout = debug.read_layout()
+    assert layout.text == "Homescreen"
+
+    features = device_handler.features()
+    assert features.initialized is True
+    assert features.recovery_mode is False

tests/upgrade_tests/test_firmware_upgrades.py

@@ -209,6 +209,7 @@ def test_upgrade_shamir_recovery(gen, from_tag, to_tag):
 
         device_handler.run(device.recover, pin_protection=False)
 
+        recovery.confirm_recovery(debug)
         recovery.select_number_of_words(debug)
         layout = recovery.enter_share(debug, MNEMONIC_SLIP39_BASIC_20_3of6[0])
         assert "2 more shares" in layout.text