tests/click: refactor recovery

tests/click_tests/recovery.py

@@ -0,0 +1,58 @@
+from .. import buttons
+
+
+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)
+
+
+def select_number_of_words(debug, num_of_words=20):
+    # confirm recovery
+    layout = debug.wait_layout()
+    assert layout.text.startswith("Recovery mode")
+    layout = debug.click(buttons.OK, wait=True)
+
+    # select number of words
+    assert "Select number of words" in layout.text
+    layout = debug.click(buttons.OK, wait=True)
+    assert layout.text == "WordSelector"
+
+    # click the number
+    word_option_offset = 6
+    word_options = (12, 18, 20, 24, 33)
+    index = word_option_offset + word_options.index(
+        num_of_words
+    )  # raises if num of words is invalid
+    coords = buttons.grid34(index % 3, index // 3)
+    layout = debug.click(coords, wait=True)
+    assert "Enter any share" in layout.text
+
+
+def enter_share(debug, share: str):
+    layout = debug.click(buttons.OK, wait=True)
+
+    assert layout.text == "Slip39Keyboard"
+    for word in share.split(" "):
+        layout = enter_word(debug, word)
+
+    return layout
+
+
+def enter_shares(debug, shares: list):
+    layout = debug.read_layout()
+    expected_text = "Enter any share"
+    remaining = len(shares)
+    for share in shares:
+        assert expected_text in layout.text
+        layout = enter_share(debug, share)
+        remaining -= 1
+        expected_text = "RecoveryHomescreen {} more".format(remaining)
+
+    assert "You have successfully recovered your wallet" in layout.text
+
+
+def finalize(debug):
+    layout = debug.click(buttons.OK, wait=True)
+    assert layout.text == "Homescreen"

tests/click_tests/test_recovery.py

@@ -18,15 +18,8 @@ import pytest
 
 from trezorlib import device, messages
 
-from .. import buttons
 from ..common import MNEMONIC_SLIP39_BASIC_20_3of6
-
+from . import recovery
-
-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.mark.skip_t1
@@ -38,37 +31,9 @@ def test_recovery(device_handler):
     assert features.initialized is False
     device_handler.run(device.recover, pin_protection=False)
 
-    # select number of words
+    recovery.select_number_of_words(debug)
-    layout = debug.wait_layout()
+    recovery.enter_shares(debug, MNEMONIC_SLIP39_BASIC_20_3of6)
-    assert layout.text.startswith("Recovery mode")
+    recovery.finalize(debug)
-    layout = debug.click(buttons.OK, wait=True)
-
-    assert "Select number of words" in layout.text
-    layout = debug.click(buttons.OK, wait=True)
-
-    assert layout.text == "WordSelector"
-    # click "20" at 2, 2
-    coords = buttons.grid34(2, 2)
-    layout = debug.click(coords, wait=True)
-    assert "Enter any share" in layout.text
-
-    expected_text = "Enter any share (20 words)"
-    remaining = len(MNEMONIC_SLIP39_BASIC_20_3of6)
-    for share in MNEMONIC_SLIP39_BASIC_20_3of6:
-        assert expected_text in layout.text
-        layout = debug.click(buttons.OK, wait=True)
-
-        assert layout.text == "Slip39Keyboard"
-        for word in share.split(" "):
-            layout = enter_word(debug, word)
-
-        remaining -= 1
-        expected_text = "RecoveryHomescreen {} more".format(remaining)
-
-    assert "You have successfully recovered your wallet" in layout.text
-    layout = debug.click(buttons.OK, wait=True)
-
-    assert layout.text == "Homescreen"
 
     assert isinstance(device_handler.result(), messages.Success)
     features = device_handler.features()