chore(tests): fix click, upgrade and persistence tests for new UI

core/src/trezor/ui/layouts/tt_v2/recovery.py

@@ -9,7 +9,7 @@ from ..common import interact
 from . import _RustLayout
 
 if TYPE_CHECKING:
-    from typing import Iterable, Callable, Any
+    from typing import Iterable, Callable
     from trezor.wire import GenericContext
 
 
@@ -42,18 +42,11 @@ async def request_word_count(ctx: GenericContext, dry_run: bool) -> int:
 async def request_word(
     ctx: GenericContext, word_index: int, word_count: int, is_slip39: bool
 ) -> str:
+    prompt = f"Type word {word_index + 1} of {word_count}:"
     if is_slip39:
-        keyboard: Any = _RustLayout(
-            trezorui2.request_bip39(
-                prompt=f"Type word {word_index + 1} of {word_count}:"
-            )
-        )
+        keyboard = _RustLayout(trezorui2.request_slip39(prompt=prompt))
     else:
-        keyboard = _RustLayout(
-            trezorui2.request_slip39(
-                prompt=f"Type word {word_index + 1} of {word_count}:"
-            )
-        )
+        keyboard = _RustLayout(trezorui2.request_bip39(prompt=prompt))
 
     word: str = await ctx.wait(keyboard)
     return word

docs/tests/click-tests.md

@@ -17,7 +17,7 @@ In the `trezor-firmware` checkout, in the root of the monorepo, install the envi
 poetry install
 ```
 
-Switch to a shell inside theenvironment:
+Switch to a shell inside the environment:
 
 ```sh
 poetry shell

tests/buttons.py

@@ -1,3 +1,4 @@
+import time
 from typing import Iterator, Tuple
 
 DISPLAY_WIDTH = 240
@@ -27,13 +28,14 @@ RESET_MINUS = (LEFT, grid(DISPLAY_HEIGHT, 5, 1))
 RESET_PLUS = (RIGHT, grid(DISPLAY_HEIGHT, 5, 1))
 
 RESET_WORD_CHECK = [
-    (MID, grid(DISPLAY_HEIGHT, 6, 3)),
-    (MID, grid(DISPLAY_HEIGHT, 6, 4)),
-    (MID, grid(DISPLAY_HEIGHT, 6, 5)),
+    (MID, grid(DISPLAY_HEIGHT, 5, 2)),
+    (MID, grid(DISPLAY_HEIGHT, 5, 3)),
+    (MID, grid(DISPLAY_HEIGHT, 5, 4)),
 ]
 
 
-BUTTON_LETTERS = ("ab", "cd", "ef", "ghij", "klm", "nopq", "rs", "tuv", "wxyz")
+BUTTON_LETTERS_BIP39 = ("abc", "def", "ghi", "jkl", "mno", "pqr", "stu", "vwx", "yz")
+BUTTON_LETTERS_SLIP39 = ("ab", "cd", "ef", "ghij", "klm", "nopq", "rs", "tuv", "wxyz")
 
 
 def grid35(x: int, y: int) -> Tuple[int, int]:
@@ -44,9 +46,39 @@ def grid34(x: int, y: int) -> Tuple[int, int]:
     return grid(DISPLAY_WIDTH, 3, x), grid(DISPLAY_HEIGHT, 4, y)
 
 
-def type_word(word: str) -> Iterator[Tuple[int, int]]:
+def _grid34_from_index(idx: int) -> Tuple[int, int]:
+    grid_x = idx % 3
+    grid_y = idx // 3 + 1  # first line is empty
+    return grid34(grid_x, grid_y)
+
+
+def type_word(word: str, is_slip39: bool = False) -> Iterator[Tuple[int, int]]:
+    if is_slip39:
+        yield from type_word_slip39(word)
+    else:
+        yield from type_word_bip39(word)
+
+
+def type_word_slip39(word: str) -> Iterator[Tuple[int, int]]:
     for l in word:
-        idx = next(i for i, letters in enumerate(BUTTON_LETTERS) if l in letters)
-        grid_x = idx % 3
-        grid_y = idx // 3 + 1  # first line is empty
-        yield grid34(grid_x, grid_y)
+        idx = next(i for i, letters in enumerate(BUTTON_LETTERS_SLIP39) if l in letters)
+        yield _grid34_from_index(idx)
+
+
+def type_word_bip39(word: str) -> Iterator[Tuple[int, int]]:
+    coords_prev: Tuple[int, int] | None = None
+    for letter in word:
+        coords, amount = letter_coords_and_amount(letter)
+        # If the button is the same as for the previous letter,
+        # waiting a second before pressing it again.
+        if coords == coords_prev:
+            time.sleep(1)
+        coords_prev = coords
+        for _ in range(amount):
+            yield coords
+
+
+def letter_coords_and_amount(letter: str) -> Tuple[Tuple[int, int], int]:
+    idx = next(i for i, letters in enumerate(BUTTON_LETTERS_BIP39) if letter in letters)
+    click_amount = BUTTON_LETTERS_BIP39[idx].index(letter) + 1
+    return _grid34_from_index(idx), click_amount

tests/click_tests/recovery.py

@@ -1,26 +1,47 @@
+from typing import TYPE_CHECKING
+
 from .. import buttons
 
+if TYPE_CHECKING:
+    from trezorlib.debuglink import DebugLink, LayoutContent
+
 
-def enter_word(debug, word):
-    word = word[:4]
-    for coords in buttons.type_word(word):
+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)
+
+    # For BIP39 - double-click on CONFIRM WORD is needed in case the word
+    # is not already typed as a whole
+    if not is_slip39 and typed_word != word:
+        debug.click(buttons.CONFIRM_WORD)
     return debug.click(buttons.CONFIRM_WORD, wait=True)
 
 
-def confirm_recovery(debug):
+def confirm_recovery(debug: "DebugLink", legacy_ui: bool = False) -> None:
     layout = debug.wait_layout()
-    assert layout.text.startswith("Recovery mode")
+    if legacy_ui:
+        layout.text.startswith("Recovery mode")
+    else:
+        assert layout.get_title() == "RECOVERY MODE"
     debug.click(buttons.OK, wait=True)
 
 
-def select_number_of_words(debug, num_of_words=20):
+def select_number_of_words(
+    debug: "DebugLink", num_of_words: int = 20, legacy_ui: bool = False
+) -> None:
     layout = debug.read_layout()
 
     # select number of words
-    assert "Select number of words" in layout.text
+    assert "Select number of words" in layout.get_content()
     layout = debug.click(buttons.OK, wait=True)
-    assert layout.text == "WordSelector"
+    if legacy_ui:
+        assert layout.text == "WordSelector"
+    else:
+        # Two title options
+        assert layout.get_title() in ("SEED CHECK", "RECOVERY MODE")
 
     # click the number
     word_option_offset = 6
@@ -30,32 +51,38 @@ def select_number_of_words(debug, num_of_words=20):
     )  # 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
+    assert "Enter any share" in layout.get_content()
 
 
-def enter_share(debug, share: str):
+def enter_share(
+    debug: "DebugLink", share: str, legacy_ui: bool = False
+) -> "LayoutContent":
     layout = debug.click(buttons.OK, wait=True)
 
-    assert layout.text == "Slip39Keyboard"
+    if legacy_ui:
+        assert layout.text == "Slip39Keyboard"
+    else:
+        assert layout.text == "< MnemonicKeyboard >"
+
     for word in share.split(" "):
-        layout = enter_word(debug, word)
+        layout = enter_word(debug, word, is_slip39=True)
 
     return layout
 
 
-def enter_shares(debug, shares: list):
+def enter_shares(debug: "DebugLink", shares: list[str]) -> None:
     layout = debug.read_layout()
     expected_text = "Enter any share"
     remaining = len(shares)
     for share in shares:
-        assert expected_text in layout.text
+        assert expected_text in layout.get_content()
         layout = enter_share(debug, share)
         remaining -= 1
-        expected_text = f"RecoveryHomescreen {remaining} more"
+        expected_text = f"{remaining} more share"
 
-    assert "You have successfully recovered your wallet" in layout.text
+    assert "You have successfully recovered your wallet" in layout.get_content()
 
 
-def finalize(debug):
+def finalize(debug: "DebugLink") -> None:
     layout = debug.click(buttons.OK, wait=True)
     assert layout.text == "Homescreen"

tests/click_tests/reset.py

@@ -1,75 +1,76 @@
+from typing import TYPE_CHECKING
+
 from shamir_mnemonic import shamir
 
 from trezorlib import messages
 
 from .. import buttons
 
+if TYPE_CHECKING:
+    from trezorlib.debuglink import DebugLink
+
 
-def confirm_wait(debug, startswith):
+def confirm_wait(debug: "DebugLink", title: str) -> None:
     layout = debug.wait_layout()
-    assert layout.text.startswith(startswith)
+    assert title.upper() in layout.get_title()
     debug.click(buttons.OK, wait=True)
 
 
-def confirm_read(debug, startswith):
+def confirm_read(debug: "DebugLink", title: str) -> None:
     layout = debug.read_layout()
-    assert layout.text.startswith(startswith)
+    if title == "Caution":
+        assert "OK, I UNDERSTAND" in layout.text
+    elif title == "Success":
+        assert any(
+            text in layout.get_content() for text in ("success", "finished", "done")
+        )
+    else:
+        assert title.upper() in layout.get_title()
     debug.click(buttons.OK, wait=True)
 
 
-def set_selection(debug, button, diff):
+def set_selection(debug: "DebugLink", button: tuple[int, int], diff: int) -> None:
     layout = debug.read_layout()
-    assert layout.text.startswith("Slip39NumInput")
+    assert "NumberInputDialog" in layout.text
     for _ in range(diff):
         debug.click(button, wait=False)
     debug.click(buttons.OK, wait=True)
 
 
-def read_words(debug, is_advanced=False):
-    def read_word(line: str):
-        return line.split()[1]
-
-    words = []
+def read_words(debug: "DebugLink", is_advanced: bool = False) -> list[str]:
+    words: list[str] = []
     layout = debug.read_layout()
     if is_advanced:
-        assert layout.text.startswith("Group")
+        assert layout.get_title().startswith("GROUP")
     else:
-        assert layout.text.startswith("Recovery share")
-
-    lines = layout.lines
-    # first screen
-    words.append(read_word(lines[3]))
-    words.append(read_word(lines[4]))
-    lines = debug.input(swipe=messages.DebugSwipeDirection.UP, wait=True).lines
-
-    # screens 2 through
-    for _ in range(4):
-        words.append(read_word(lines[1]))
-        words.append(read_word(lines[2]))
-        words.append(read_word(lines[3]))
-        words.append(read_word(lines[4]))
-        lines = debug.input(swipe=messages.DebugSwipeDirection.UP, wait=True).lines
-
-    # final screen
-    words.append(read_word(lines[1]))
-    words.append(read_word(lines[2]))
+        assert layout.get_title().startswith("RECOVERY SHARE #")
+
+    # Swiping through all the page and loading the words
+    for _ in range(layout.get_page_count() - 1):
+        words.extend(layout.get_seed_words())
+        layout = debug.input(swipe=messages.DebugSwipeDirection.UP, wait=True)
+    words.extend(layout.get_seed_words())
+
     debug.press_yes()
 
     return words
 
 
-def confirm_words(debug, words):
+def confirm_words(debug: "DebugLink", words: list[str]) -> None:
     layout = debug.wait_layout()
     assert "Select word" in layout.text
     for _ in range(3):
         # "Select word 3 of 20"
         #              ^
-        word_pos = int(layout.lines[1].split()[2])
-        button_pos = layout.lines.index(words[word_pos - 1]) - 2
+        word_pos = int(layout.get_content().split()[2])
+        # Unifying both the buttons and words to lowercase
+        btn_texts = [text.lower() for text in layout.get_button_texts()]
+        wanted_word = words[word_pos - 1].lower()
+        button_pos = btn_texts.index(wanted_word)
         layout = debug.click(buttons.RESET_WORD_CHECK[button_pos], wait=True)
 
 
-def validate_mnemonics(mnemonics, expected_ems):
+def validate_mnemonics(mnemonics: list[str], expected_ems: bytes) -> None:
     # We expect these combinations to recreate the secret properly
     # In case of click tests the mnemonics are always XofX so no need for combinations
     groups = shamir.decode_mnemonics(mnemonics)

tests/click_tests/test_autolock.py

@@ -15,6 +15,7 @@
 # If not, see <https://www.gnu.org/licenses/lgpl-3.0.html>.
 
 import time
+from typing import TYPE_CHECKING
 
 import pytest
 
@@ -25,6 +26,9 @@ from .. import buttons, common
 from ..tx_cache import TxCache
 from . import recovery
 
+if TYPE_CHECKING:
+    from ..device_handler import BackgroundDeviceHandler
+
 TX_CACHE = TxCache("Bitcoin")
 
 TXHASH_d5f65e = bytes.fromhex(
@@ -36,17 +40,20 @@ PIN4 = "1234"
 WORDS_20 = buttons.grid34(2, 2)
 
 
-def set_autolock_delay(device_handler, delay_ms):
+def set_autolock_delay(device_handler: "BackgroundDeviceHandler", delay_ms: int):
     debug = device_handler.debuglink()
 
     device_handler.run(device.apply_settings, auto_lock_delay_ms=delay_ms)
 
     layout = debug.wait_layout()
-    assert layout.text == "PinDialog"
+    assert layout.text == "< PinKeyboard >"
     debug.input("1234")
 
     layout = debug.wait_layout()
-    assert f"auto-lock your device after  {delay_ms // 1000} seconds" in layout.text
+    assert (
+        f"auto-lock your device after {delay_ms // 1000} seconds"
+        in layout.get_content()
+    )
     debug.click(buttons.OK)
 
     layout = debug.wait_layout()
@@ -55,7 +62,7 @@ def set_autolock_delay(device_handler, delay_ms):
 
 
 @pytest.mark.setup_client(pin=PIN4)
-def test_autolock_interrupts_signing(device_handler):
+def test_autolock_interrupts_signing(device_handler: "BackgroundDeviceHandler"):
     set_autolock_delay(device_handler, 10_000)
 
     debug = device_handler.debuglink()
@@ -76,10 +83,13 @@ def test_autolock_interrupts_signing(device_handler):
     device_handler.run(btc.sign_tx, "Bitcoin", [inp1], [out1], prev_txes=TX_CACHE)
 
     layout = debug.wait_layout()
-    assert "1MJ2tj2ThBE62zXbBYA5ZaN3fdve5CPAz1" in layout.text.replace(" ", "")
+    assert "1MJ2tj2ThBE62zXbBYA5ZaN3fdve5CPAz1" in layout.get_content().replace(" ", "")
+
+    debug.click(buttons.OK, wait=True)
+    debug.click(buttons.OK, wait=True)
 
     layout = debug.click(buttons.OK, wait=True)
-    assert "Total amount:  0.0039 BTC" in layout.text
+    assert "Total amount: 0.0039 BTC" in layout.get_content()
 
     # wait for autolock to kick in
     time.sleep(10.1)
@@ -89,7 +99,7 @@ def test_autolock_interrupts_signing(device_handler):
 
 @pytest.mark.xfail(reason="depends on #922")
 @pytest.mark.setup_client(pin=PIN4, passphrase=True)
-def test_autolock_passphrase_keyboard(device_handler):
+def test_autolock_passphrase_keyboard(device_handler: "BackgroundDeviceHandler"):
     set_autolock_delay(device_handler, 10_000)
     debug = device_handler.debuglink()
 
@@ -109,7 +119,7 @@ def test_autolock_passphrase_keyboard(device_handler):
 
 
 @pytest.mark.setup_client(pin=PIN4)
-def test_dryrun_locks_at_number_of_words(device_handler):
+def test_dryrun_locks_at_number_of_words(device_handler: "BackgroundDeviceHandler"):
     set_autolock_delay(device_handler, 10_000)
     debug = device_handler.debuglink()
 
@@ -117,11 +127,11 @@ def test_dryrun_locks_at_number_of_words(device_handler):
 
     # unlock
     layout = debug.wait_layout()
-    assert "Do you really want to check the recovery seed?" in layout.text
+    assert "Do you really want to check the recovery seed?" in layout.get_content()
     layout = debug.click(buttons.OK, wait=True)
-    assert layout.text == "PinDialog"
+    assert layout.text == "< PinKeyboard >"
     layout = debug.input(PIN4, wait=True)
-    assert "Select number of words " in layout.text
+    assert "Select number of words " in layout.get_content()
 
     # wait for autolock to trigger
     time.sleep(10.1)
@@ -132,15 +142,15 @@ def test_dryrun_locks_at_number_of_words(device_handler):
 
     # unlock
     layout = debug.click(buttons.OK, wait=True)
-    assert layout.text == "PinDialog"
+    assert layout.text == "< PinKeyboard >"
     layout = debug.input(PIN4, wait=True)
 
     # we are back at homescreen
-    assert "Select number of words" in layout.text
+    assert "Select number of words" in layout.get_content()
 
 
 @pytest.mark.setup_client(pin=PIN4)
-def test_dryrun_locks_at_word_entry(device_handler):
+def test_dryrun_locks_at_word_entry(device_handler: "BackgroundDeviceHandler"):
     set_autolock_delay(device_handler, 10_000)
     debug = device_handler.debuglink()
 
@@ -148,9 +158,9 @@ def test_dryrun_locks_at_word_entry(device_handler):
 
     # unlock
     layout = debug.wait_layout()
-    assert "Do you really want to check the recovery seed?" in layout.text
+    assert "Do you really want to check the recovery seed?" in layout.get_content()
     layout = debug.click(buttons.OK, wait=True)
-    assert layout.text == "PinDialog"
+    assert layout.text == "< PinKeyboard >"
     layout = debug.input(PIN4, wait=True)
 
     # select 20 words
@@ -158,7 +168,7 @@ def test_dryrun_locks_at_word_entry(device_handler):
 
     layout = debug.click(buttons.OK, wait=True)
     # make sure keyboard locks
-    assert layout.text == "Slip39Keyboard"
+    assert layout.text == "< MnemonicKeyboard >"
     time.sleep(10.1)
     layout = debug.wait_layout()
     assert layout.text == "Lockscreen"
@@ -167,7 +177,7 @@ def test_dryrun_locks_at_word_entry(device_handler):
 
 
 @pytest.mark.setup_client(pin=PIN4)
-def test_dryrun_enter_word_slowly(device_handler):
+def test_dryrun_enter_word_slowly(device_handler: "BackgroundDeviceHandler"):
     set_autolock_delay(device_handler, 10_000)
     debug = device_handler.debuglink()
 
@@ -175,9 +185,9 @@ def test_dryrun_enter_word_slowly(device_handler):
 
     # unlock
     layout = debug.wait_layout()
-    assert "Do you really want to check the recovery seed?" in layout.text
+    assert "Do you really want to check the recovery seed?" in layout.get_content()
     layout = debug.click(buttons.OK, wait=True)
-    assert layout.text == "PinDialog"
+    assert layout.text == "< PinKeyboard >"
     layout = debug.input(PIN4, wait=True)
 
     # select 20 words
@@ -185,11 +195,11 @@ def test_dryrun_enter_word_slowly(device_handler):
 
     layout = debug.click(buttons.OK, wait=True)
     # type the word OCEAN slowly
-    assert layout.text == "Slip39Keyboard"
-    for coords in buttons.type_word("ocea"):
+    assert layout.text == "< MnemonicKeyboard >"
+    for coords in buttons.type_word("ocea", is_slip39=True):
         time.sleep(9)
         debug.click(coords)
     layout = debug.click(buttons.CONFIRM_WORD, wait=True)
     # should not have locked, even though we took 9 seconds to type each letter
-    assert layout.text == "Slip39Keyboard"
+    assert layout.text == "< MnemonicKeyboard >"
     device_handler.kill_task()

tests/click_tests/test_lock.py

@@ -15,19 +15,24 @@
 # If not, see <https://www.gnu.org/licenses/lgpl-3.0.html>.
 
 import time
+from typing import TYPE_CHECKING
 
 import pytest
 
 from .. import buttons, common
 
+if TYPE_CHECKING:
+    from ..device_handler import BackgroundDeviceHandler
+
+
 PIN4 = "1234"
 
 
 @pytest.mark.setup_client(pin=PIN4)
-def test_hold_to_lock(device_handler):
+def test_hold_to_lock(device_handler: "BackgroundDeviceHandler"):
     debug = device_handler.debuglink()
 
-    def hold(duration, wait=True):
+    def hold(duration: int, wait: bool = True) -> None:
         debug.input(x=13, y=37, hold_ms=duration, wait=wait)
         time.sleep(duration / 1000 + 0.5)
 
@@ -36,7 +41,7 @@ def test_hold_to_lock(device_handler):
     # unlock with message
     device_handler.run(common.get_test_address)
     layout = debug.wait_layout()
-    assert layout.text == "PinDialog"
+    assert layout.text == "< PinKeyboard >"
     debug.input("1234", wait=True)
     assert device_handler.result()
 
@@ -52,7 +57,7 @@ def test_hold_to_lock(device_handler):
 
     # unlock by touching
     layout = debug.click(buttons.INFO, wait=True)
-    assert layout.text == "PinDialog"
+    assert layout.text == "< PinKeyboard >"
     debug.input("1234", wait=True)
 
     assert device_handler.features().unlocked is True

tests/click_tests/test_recovery.py

@@ -14,6 +14,8 @@
 # You should have received a copy of the License along with this library.
 # If not, see <https://www.gnu.org/licenses/lgpl-3.0.html>.
 
+from typing import TYPE_CHECKING
+
 import pytest
 
 from trezorlib import device, messages
@@ -21,10 +23,13 @@ from trezorlib import device, messages
 from ..common import MNEMONIC_SLIP39_BASIC_20_3of6
 from . import recovery
 
+if TYPE_CHECKING:
+    from ..device_handler import BackgroundDeviceHandler
+
 
 @pytest.mark.skip_t1
 @pytest.mark.setup_client(uninitialized=True)
-def test_recovery(device_handler):
+def test_recovery(device_handler: "BackgroundDeviceHandler"):
     features = device_handler.features()
     debug = device_handler.debuglink()
 

tests/click_tests/test_reset_slip39_advanced.py

@@ -14,6 +14,7 @@
 # You should have received a copy of the License along with this library.
 # If not, see <https://www.gnu.org/licenses/lgpl-3.0.html>.
 
+from typing import TYPE_CHECKING
 from unittest import mock
 
 import pytest
@@ -24,6 +25,10 @@ from .. import buttons
 from ..common import generate_entropy
 from . import reset
 
+if TYPE_CHECKING:
+    from ..device_handler import BackgroundDeviceHandler
+
+
 EXTERNAL_ENTROPY = b"zlutoucky kun upel divoke ody" * 2
 
 with_mock_urandom = mock.patch("os.urandom", mock.Mock(return_value=EXTERNAL_ENTROPY))
@@ -32,7 +37,9 @@ with_mock_urandom = mock.patch("os.urandom", mock.Mock(return_value=EXTERNAL_ENT
 @pytest.mark.skip_t1
 @pytest.mark.setup_client(uninitialized=True)
 @with_mock_urandom
-def test_reset_slip39_advanced_2of2groups_2of2shares(device_handler):
+def test_reset_slip39_advanced_2of2groups_2of2shares(
+    device_handler: "BackgroundDeviceHandler",
+):
     features = device_handler.features()
     debug = device_handler.debuglink()
 
@@ -76,7 +83,7 @@ def test_reset_slip39_advanced_2of2groups_2of2shares(device_handler):
     # confirm backup warning
     reset.confirm_read(debug, "Caution")
 
-    all_words = []
+    all_words: list[str] = []
     for _ in range(2):
         for _ in range(2):
             # read words
@@ -95,6 +102,7 @@ def test_reset_slip39_advanced_2of2groups_2of2shares(device_handler):
 
     # generate secret locally
     internal_entropy = debug.state().reset_entropy
+    assert internal_entropy is not None
     secret = generate_entropy(128, internal_entropy, EXTERNAL_ENTROPY)
 
     # validate that all combinations will result in the correct master secret
@@ -114,7 +122,9 @@ def test_reset_slip39_advanced_2of2groups_2of2shares(device_handler):
 @pytest.mark.setup_client(uninitialized=True)
 @pytest.mark.slow
 @with_mock_urandom
-def test_reset_slip39_advanced_16of16groups_16of16shares(device_handler):
+def test_reset_slip39_advanced_16of16groups_16of16shares(
+    device_handler: "BackgroundDeviceHandler",
+):
     features = device_handler.features()
     debug = device_handler.debuglink()
 
@@ -158,7 +168,7 @@ def test_reset_slip39_advanced_16of16groups_16of16shares(device_handler):
     # confirm backup warning
     reset.confirm_read(debug, "Caution")
 
-    all_words = []
+    all_words: list[str] = []
     for _ in range(16):
         for _ in range(16):
             # read words
@@ -177,6 +187,7 @@ def test_reset_slip39_advanced_16of16groups_16of16shares(device_handler):
 
     # generate secret locally
     internal_entropy = debug.state().reset_entropy
+    assert internal_entropy is not None
     secret = generate_entropy(128, internal_entropy, EXTERNAL_ENTROPY)
 
     # validate that all combinations will result in the correct master secret

tests/click_tests/test_reset_slip39_basic.py

@@ -14,6 +14,7 @@
 # You should have received a copy of the License along with this library.
 # If not, see <https://www.gnu.org/licenses/lgpl-3.0.html>.
 
+from typing import TYPE_CHECKING
 from unittest import mock
 
 import pytest
@@ -24,12 +25,16 @@ from .. import buttons
 from ..common import generate_entropy
 from . import reset
 
+if TYPE_CHECKING:
+    from ..device_handler import BackgroundDeviceHandler
+
+
 EXTERNAL_ENTROPY = b"zlutoucky kun upel divoke ody" * 2
 
 
 @pytest.mark.skip_t1
 @pytest.mark.setup_client(uninitialized=True)
-def test_reset_slip39_basic_1of1(device_handler):
+def test_reset_slip39_basic_1of1(device_handler: "BackgroundDeviceHandler"):
     features = device_handler.features()
     debug = device_handler.debuglink()
 
@@ -84,6 +89,7 @@ def test_reset_slip39_basic_1of1(device_handler):
 
         # generate secret locally
         internal_entropy = debug.state().reset_entropy
+        assert internal_entropy is not None
         secret = generate_entropy(128, internal_entropy, EXTERNAL_ENTROPY)
 
         # validate that all combinations will result in the correct master secret
@@ -101,7 +107,7 @@ def test_reset_slip39_basic_1of1(device_handler):
 
 @pytest.mark.skip_t1
 @pytest.mark.setup_client(uninitialized=True)
-def test_reset_slip39_basic_16of16(device_handler):
+def test_reset_slip39_basic_16of16(device_handler: "BackgroundDeviceHandler"):
     features = device_handler.features()
     debug = device_handler.debuglink()
 
@@ -142,7 +148,7 @@ def test_reset_slip39_basic_16of16(device_handler):
         # confirm backup warning
         reset.confirm_read(debug, "Caution")
 
-        all_words = []
+        all_words: list[str] = []
         for _ in range(16):
             # read words
             words = reset.read_words(debug)
@@ -160,6 +166,7 @@ def test_reset_slip39_basic_16of16(device_handler):
 
         # generate secret locally
         internal_entropy = debug.state().reset_entropy
+        assert internal_entropy is not None
         secret = generate_entropy(128, internal_entropy, EXTERNAL_ENTROPY)
 
         # validate that all combinations will result in the correct master secret

tests/device_handler.py

@@ -23,7 +23,7 @@ class NullUI:
         raise NotImplementedError("NullUI should not be used with T1")
 
     @staticmethod
-    def get_passphrase(available_on_device=False):
+    def get_passphrase(available_on_device: bool = False):
         if available_on_device:
             return PASSPHRASE_ON_DEVICE
         else:
@@ -42,7 +42,7 @@ class BackgroundDeviceHandler:
         self.client.ui = NullUI  # type: ignore [NullUI is OK UI]
         self.client.watch_layout(True)
 
-    def run(self, function, *args, **kwargs):
+    def run(self, function, *args, **kwargs) -> None:
         if self.task is not None:
             raise RuntimeError("Wait for previous task first")
         self.task = self._pool.submit(function, self.client, *args, **kwargs)
@@ -60,7 +60,7 @@ class BackgroundDeviceHandler:
                 pass
         self.task = None
 
-    def restart(self, emulator: "Emulator"):
+    def restart(self, emulator: "Emulator") -> None:
         # TODO handle actual restart as well
         self.kill_task()
         emulator.restart()

tests/persistence_tests/test_safety_checks.py

@@ -1,3 +1,5 @@
+from typing import Iterator
+
 import pytest
 
 from trezorlib import debuglink, device
@@ -9,7 +11,7 @@ from ..upgrade_tests import core_only
 
 
 @pytest.fixture
-def emulator() -> Emulator:
+def emulator() -> Iterator[Emulator]:
     with EmulatorWrapper("core") as emu:
         yield emu
 
@@ -26,7 +28,6 @@ def emulator() -> Emulator:
 def test_safety_checks_level_after_reboot(
     emulator: Emulator, set_level: SafetyCheckLevel, after_level: SafetyCheckLevel
 ):
-    assert emulator.client is not None
     device.wipe(emulator.client)
     debuglink.load_device(
         emulator.client,

tests/persistence_tests/test_shamir_persistence.py

@@ -14,6 +14,8 @@
 # You should have received a copy of the License along with this library.
 # If not, see <https://www.gnu.org/licenses/lgpl-3.0.html>.
 
+from typing import Iterator
+
 import pytest
 
 from trezorlib import device
@@ -28,7 +30,7 @@ from ..upgrade_tests import core_only
 
 
 @pytest.fixture
-def emulator() -> Emulator:
+def emulator() -> Iterator[Emulator]:
     with EmulatorWrapper("core") as emu:
         yield emu
 
@@ -48,7 +50,7 @@ def test_abort(emulator: Emulator):
 
     device_handler.run(device.recover, pin_protection=False)
     layout = debug.wait_layout()
-    assert layout.text.startswith("Recovery mode")
+    assert layout.get_title() == "RECOVERY MODE"
 
     layout = debug.click(buttons.OK, wait=True)
     assert "Select number of words" in layout.text
@@ -64,7 +66,7 @@ def test_abort(emulator: Emulator):
     assert "Select number of words" in layout.text
     layout = debug.click(buttons.CANCEL, wait=True)
 
-    assert layout.text.startswith("Abort recovery")
+    assert layout.get_title() == "ABORT RECOVERY"
     layout = debug.click(buttons.OK, wait=True)
 
     assert layout.text == "Homescreen"
@@ -137,7 +139,7 @@ def test_recovery_on_old_wallet(emulator: Emulator):
     assert "Enter any share" in layout.text
     debug.press_yes()
     layout = debug.wait_layout()
-    assert layout.text == "Slip39Keyboard"
+    assert layout.text == "< MnemonicKeyboard >"
 
     # enter first word
     debug.input(words[0])
@@ -149,7 +151,7 @@ def test_recovery_on_old_wallet(emulator: Emulator):
 
     # try entering remaining 19 words
     for word in words[1:]:
-        assert layout.text == "Slip39Keyboard"
+        assert layout.text == "< MnemonicKeyboard >"
         debug.input(word)
         layout = debug.wait_layout()
 
@@ -179,10 +181,10 @@ def test_recovery_multiple_resets(emulator: Emulator):
             assert expected_text in layout.text
             layout = recovery.enter_share(debug, share)
             remaining -= 1
-            expected_text = "Success You have entered"
+            expected_text = "You have entered"
             debug = _restart(device_handler, emulator)
 
-        assert "You have successfully recovered your wallet" in layout.text
+        assert "You have successfully recovered your wallet" in layout.get_content()
 
     device_handler = BackgroundDeviceHandler(emulator.client)
     debug = device_handler.debuglink()

tests/persistence_tests/test_wipe_code.py

@@ -50,7 +50,6 @@ def setup_device_core(client: Client, pin: str, wipe_code: str) -> None:
 @core_only
 def test_wipe_code_activate_core():
     with EmulatorWrapper("core") as emu:
-        assert emu.client is not None
         # set up device
         setup_device_core(emu.client, PIN, WIPE_CODE)
 
@@ -82,7 +81,6 @@ def test_wipe_code_activate_core():
 @legacy_only
 def test_wipe_code_activate_legacy():
     with EmulatorWrapper("legacy") as emu:
-        assert emu.client is not None
         # set up device
         setup_device_legacy(emu.client, PIN, WIPE_CODE)
 

tests/upgrade_tests/__init__.py

@@ -47,7 +47,7 @@ core_only = pytest.mark.skipif(
 
 
 def for_all(
-    *args,
+    *args: str,
     legacy_minimum_version: Tuple[int, int, int] = (1, 0, 0),
     core_minimum_version: Tuple[int, int, int] = (2, 0, 0)
 ) -> "MarkDecorator":
@@ -71,7 +71,7 @@ def for_all(
     # If any gens were selected, use them. If none, select all.
     enabled_gens = SELECTED_GENS or args
 
-    all_params = []
+    all_params: list[tuple[str, str | None]] = []
     for gen in args:
         if gen == "legacy":
             minimum_version = legacy_minimum_version

tests/upgrade_tests/test_firmware_upgrades.py

@@ -15,7 +15,7 @@
 # If not, see <https://www.gnu.org/licenses/lgpl-3.0.html>.
 
 import dataclasses
-from typing import List
+from typing import TYPE_CHECKING, List, Optional
 
 import pytest
 
@@ -29,6 +29,9 @@ from ..device_handler import BackgroundDeviceHandler
 from ..emulators import ALL_TAGS, EmulatorWrapper
 from . import for_all, for_tags
 
+if TYPE_CHECKING:
+    from trezorlib.debuglink import TrezorClientDebugLink as Client
+
 models.TREZOR_ONE = dataclasses.replace(models.TREZOR_ONE, minimum_version=(1, 0, 0))
 models.TREZOR_T = dataclasses.replace(models.TREZOR_T, minimum_version=(2, 0, 0))
 models.TREZORS = {models.TREZOR_ONE, models.TREZOR_T}
@@ -44,8 +47,8 @@ STRENGTH = 128
 
 
 @for_all()
-def test_upgrade_load(gen: str, tag: str):
-    def asserts(client):
+def test_upgrade_load(gen: str, tag: str) -> None:
+    def asserts(client: "Client"):
         assert not client.features.pin_protection
         assert not client.features.passphrase_protection
         assert client.features.initialized
@@ -72,10 +75,10 @@ def test_upgrade_load(gen: str, tag: str):
 
 
 @for_all("legacy")
-def test_upgrade_load_pin(gen: str, tag: str):
+def test_upgrade_load_pin(gen: str, tag: str) -> None:
     PIN = "1234"
 
-    def asserts(client):
+    def asserts(client: "Client") -> None:
         assert client.features.pin_protection
         assert not client.features.passphrase_protection
         assert client.features.initialized
@@ -117,7 +120,7 @@ def test_upgrade_load_pin(gen: str, tag: str):
 def test_storage_upgrade_progressive(gen: str, tags: List[str]):
     PIN = "1234"
 
-    def asserts(client):
+    def asserts(client: "Client") -> None:
         assert client.features.pin_protection
         assert not client.features.passphrase_protection
         assert client.features.initialized
@@ -153,7 +156,7 @@ def test_upgrade_wipe_code(gen: str, tag: str):
     PIN = "1234"
     WIPE_CODE = "4321"
 
-    def asserts(client):
+    def asserts(client: "Client"):
         assert client.features.pin_protection
         assert not client.features.passphrase_protection
         assert client.features.initialized
@@ -195,7 +198,7 @@ def test_upgrade_wipe_code(gen: str, tag: str):
 
 @for_all("legacy")
 def test_upgrade_reset(gen: str, tag: str):
-    def asserts(client):
+    def asserts(client: "Client"):
         assert not client.features.pin_protection
         assert not client.features.passphrase_protection
         assert client.features.initialized
@@ -228,7 +231,7 @@ def test_upgrade_reset(gen: str, tag: str):
 
 @for_all()
 def test_upgrade_reset_skip_backup(gen: str, tag: str):
-    def asserts(client):
+    def asserts(client: "Client"):
         assert not client.features.pin_protection
         assert not client.features.passphrase_protection
         assert client.features.initialized
@@ -262,7 +265,7 @@ def test_upgrade_reset_skip_backup(gen: str, tag: str):
 
 @for_all(legacy_minimum_version=(1, 7, 2))
 def test_upgrade_reset_no_backup(gen: str, tag: str):
-    def asserts(client):
+    def asserts(client: "Client"):
         assert not client.features.pin_protection
         assert not client.features.passphrase_protection
         assert client.features.initialized
@@ -296,7 +299,7 @@ def test_upgrade_reset_no_backup(gen: str, tag: str):
 
 # Although Shamir was introduced in 2.1.2 already, the debug instrumentation was not present until 2.1.9.
 @for_all("core", core_minimum_version=(2, 1, 9))
-def test_upgrade_shamir_recovery(gen: str, tag: str):
+def test_upgrade_shamir_recovery(gen: str, tag: Optional[str]):
     with EmulatorWrapper(gen, tag) as emu, BackgroundDeviceHandler(
         emu.client
     ) as device_handler:
@@ -306,9 +309,14 @@ def test_upgrade_shamir_recovery(gen: str, tag: str):
 
         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])
+        # Flow is different for old UI and new UI
+        legacy_ui = emu.client.version < (2, 5, 4)
+
+        recovery.confirm_recovery(debug, legacy_ui=legacy_ui)
+        recovery.select_number_of_words(debug, legacy_ui=legacy_ui)
+        layout = recovery.enter_share(
+            debug, MNEMONIC_SLIP39_BASIC_20_3of6[0], legacy_ui=legacy_ui
+        )
         assert "2 more shares" in layout.text
 
         device_id = emu.client.features.device_id
@@ -331,6 +339,7 @@ def test_upgrade_shamir_recovery(gen: str, tag: str):
 
         # Check the result
         state = debug.state()
+        assert state.mnemonic_secret is not None
         assert state.mnemonic_secret.hex() == MNEMONIC_SLIP39_BASIC_20_3of6_SECRET
         assert state.mnemonic_type == BackupType.Slip39_Basic
 

tests/upgrade_tests/test_passphrase_consistency.py

@@ -14,6 +14,8 @@
 # You should have received a copy of the License along with this library.
 # If not, see <https://www.gnu.org/licenses/lgpl-3.0.html>.
 
+from typing import Iterator
+
 import pytest
 
 from trezorlib import btc, device, mapping, messages, models, protobuf
@@ -41,9 +43,8 @@ mapping.DEFAULT_MAPPING.register(ApplySettingsCompat)
 
 
 @pytest.fixture
-def emulator(gen: str, tag: str) -> Emulator:
+def emulator(gen: str, tag: str) -> Iterator[Emulator]:
     with EmulatorWrapper(gen, tag) as emu:
-        assert emu.client is not None
         # set up a passphrase-protected device
         device.reset(
             emu.client,
@@ -64,7 +65,6 @@ def emulator(gen: str, tag: str) -> Emulator:
 )
 def test_passphrase_works(emulator: Emulator):
     """Check that passphrase handling in trezorlib works correctly in all versions."""
-    assert emulator.client is not None
     if emulator.client.features.model == "T" and emulator.client.version < (2, 3, 0):
         expected_responses = [
             messages.PassphraseRequest,
@@ -102,7 +102,6 @@ def test_init_device(emulator: Emulator):
     """Check that passphrase caching and session_id retaining works correctly across
     supported versions.
     """
-    assert emulator.client is not None
     if emulator.client.features.model == "T" and emulator.client.version < (2, 3, 0):
         expected_responses = [
             messages.PassphraseRequest,