chore(tests): use special recovery functions for upgrade tests

[no changelog]

tests/upgrade_tests/recovery_old.py

@@ -0,0 +1,46 @@
+from typing import TYPE_CHECKING
+
+from .. import buttons
+
+if TYPE_CHECKING:
+    from trezorlib.debuglink import DebugLink, LayoutContent
+
+
+def _enter_word(
+    debug: "DebugLink", word: str, is_slip39: bool = False
+) -> "LayoutContent":
+    typed_word = word[:4]
+    for coords in buttons.type_word(typed_word, is_slip39=is_slip39):
+        debug.click(coords)
+
+    return debug.click(buttons.CONFIRM_WORD, wait=True)
+
+
+def confirm_recovery(debug: "DebugLink") -> None:
+    debug.click(buttons.OK, wait=True)
+
+
+def select_number_of_words(
+    debug: "DebugLink", num_of_words: int = 20, wait: bool = True
+) -> None:
+    if wait:
+        debug.wait_layout()
+    debug.click(buttons.OK, wait=True)
+
+    # 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)
+    debug.click(coords, wait=True)
+
+
+def enter_share(
+    debug: "DebugLink", share: str, is_first: bool = True
+) -> "LayoutContent":
+    layout = debug.click(buttons.OK, wait=True)
+    for word in share.split(" "):
+        layout = _enter_word(debug, word, is_slip39=True)
+    return layout

tests/upgrade_tests/test_firmware_upgrades.py

@@ -23,11 +23,10 @@ from trezorlib import btc, debuglink, device, exceptions, fido, models
 from trezorlib.messages import BackupType
 from trezorlib.tools import H_
 
-from ..click_tests import recovery
 from ..common import MNEMONIC_SLIP39_BASIC_20_3of6, MNEMONIC_SLIP39_BASIC_20_3of6_SECRET
 from ..device_handler import BackgroundDeviceHandler
 from ..emulators import ALL_TAGS, EmulatorWrapper
-from . import for_all, for_tags
+from . import for_all, for_tags, recovery_old
 
 if TYPE_CHECKING:
     from trezorlib.debuglink import TrezorClientDebugLink as Client
@@ -309,9 +308,9 @@ def test_upgrade_shamir_recovery(gen: str, tag: Optional[str]):
 
         device_handler.run(device.recover, pin_protection=False)
 
-        recovery.confirm_recovery(debug)
+        recovery_old.confirm_recovery(debug)
-        recovery.select_number_of_words(debug, wait=not debug.legacy_debug)
+        recovery_old.select_number_of_words(debug, wait=not debug.legacy_debug)
-        layout = recovery.enter_share(debug, MNEMONIC_SLIP39_BASIC_20_3of6[0])
+        layout = recovery_old.enter_share(debug, MNEMONIC_SLIP39_BASIC_20_3of6[0])
         if not debug.legacy_ui and not debug.legacy_debug:
             assert (
                 "1 of 3 shares entered" in layout.text_content()
@@ -329,14 +328,14 @@ def test_upgrade_shamir_recovery(gen: str, tag: Optional[str]):
         emu.client.watch_layout(True)
 
         # second share
-        layout = recovery.enter_share(debug, MNEMONIC_SLIP39_BASIC_20_3of6[2])
+        layout = recovery_old.enter_share(debug, MNEMONIC_SLIP39_BASIC_20_3of6[2])
         assert (
             "2 of 3 shares entered" in layout.text_content()
             or "1 more share" in layout.text_content()
         )
 
         # last one
-        layout = recovery.enter_share(debug, MNEMONIC_SLIP39_BASIC_20_3of6[1])
+        layout = recovery_old.enter_share(debug, MNEMONIC_SLIP39_BASIC_20_3of6[1])
         assert (
             "Wallet recovered successfully" in layout.text_content()
             or "finished recovering" in layout.text_content()