chore(tests): add type hints to some non-device tests

tests/persistence_tests/test_safety_checks.py

@@ -4,12 +4,12 @@ from trezorlib import debuglink, device
 from trezorlib.messages import SafetyCheckLevel
 
 from ..common import MNEMONIC12
-from ..emulators import EmulatorWrapper
+from ..emulators import Emulator, EmulatorWrapper
 from ..upgrade_tests import core_only
 
 
 @pytest.fixture
-def emulator():
+def emulator() -> Emulator:
     with EmulatorWrapper("core") as emu:
         yield emu
 
@@ -23,7 +23,10 @@ def emulator():
         (SafetyCheckLevel.PromptAlways, SafetyCheckLevel.PromptAlways),
     ],
 )
-def test_safety_checks_level_after_reboot(emulator, set_level, after_level):
+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

@@ -17,28 +17,29 @@
 import pytest
 
 from trezorlib import device
+from trezorlib.debuglink import DebugLink
 
 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 ..emulators import Emulator, EmulatorWrapper
 from ..upgrade_tests import core_only
 
 
 @pytest.fixture
-def emulator():
+def emulator() -> Emulator:
     with EmulatorWrapper("core") as emu:
         yield emu
 
 
-def _restart(device_handler, emulator):
+def _restart(device_handler: BackgroundDeviceHandler, emulator: Emulator):
     device_handler.restart(emulator)
     return device_handler.debuglink()
 
 
 @core_only
-def test_abort(emulator):
+def test_abort(emulator: Emulator):
     device_handler = BackgroundDeviceHandler(emulator.client)
     debug = device_handler.debuglink()
     features = device_handler.features()
@@ -72,7 +73,7 @@ def test_abort(emulator):
 
 
 @core_only
-def test_recovery_single_reset(emulator):
+def test_recovery_single_reset(emulator: Emulator):
     device_handler = BackgroundDeviceHandler(emulator.client)
     debug = device_handler.debuglink()
     features = device_handler.features()
@@ -100,7 +101,7 @@ def test_recovery_single_reset(emulator):
 
 
 @core_only
-def test_recovery_on_old_wallet(emulator):
+def test_recovery_on_old_wallet(emulator: Emulator):
     """Check that the recovery workflow started on a disconnected device can survive
     handling by the old Wallet.
 
@@ -168,8 +169,8 @@ def test_recovery_on_old_wallet(emulator):
 
 
 @core_only
-def test_recovery_multiple_resets(emulator):
-    def enter_shares_with_restarts(debug):
+def test_recovery_multiple_resets(emulator: Emulator):
+    def enter_shares_with_restarts(debug: DebugLink) -> None:
         shares = MNEMONIC_SLIP39_ADVANCED_20
         layout = debug.read_layout()
         expected_text = "Enter any share"

tests/persistence_tests/test_wipe_code.py

@@ -8,7 +8,9 @@ PIN = "1234"
 WIPE_CODE = "9876"
 
 
-def setup_device_legacy(client, pin, wipe_code):
+def setup_device_legacy(
+    client: debuglink.TrezorClientDebugLink, pin: str, wipe_code: str
+) -> None:
     device.wipe(client)
     debuglink.load_device(
         client, MNEMONIC12, pin, passphrase_protection=False, label="WIPECODE"
@@ -19,7 +21,9 @@ def setup_device_legacy(client, pin, wipe_code):
         device.change_wipe_code(client)
 
 
-def setup_device_core(client, pin, wipe_code):
+def setup_device_core(
+    client: debuglink.TrezorClientDebugLink, pin: str, wipe_code: str
+) -> None:
     device.wipe(client)
     debuglink.load_device(
         client, MNEMONIC12, pin, passphrase_protection=False, label="WIPECODE"
@@ -49,6 +53,7 @@ def setup_device_core(client, pin, wipe_code):
 @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)
 
@@ -80,6 +85,7 @@ 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/test_firmware_upgrades.py

@@ -15,6 +15,7 @@
 # If not, see <https://www.gnu.org/licenses/lgpl-3.0.html>.
 
 import dataclasses
+from typing import List
 
 import pytest
 
@@ -43,7 +44,7 @@ STRENGTH = 128
 
 
 @for_all()
-def test_upgrade_load(gen, tag):
+def test_upgrade_load(gen: str, tag: str):
     def asserts(client):
         assert not client.features.pin_protection
         assert not client.features.passphrase_protection
@@ -71,7 +72,7 @@ def test_upgrade_load(gen, tag):
 
 
 @for_all("legacy")
-def test_upgrade_load_pin(gen, tag):
+def test_upgrade_load_pin(gen: str, tag: str):
     PIN = "1234"
 
     def asserts(client):
@@ -113,7 +114,7 @@ def test_upgrade_load_pin(gen, tag):
     ("legacy", ["v1.7.0", "v1.9.0"]),
     ("legacy", ["v1.8.0", "v1.9.0"]),
 )
-def test_storage_upgrade_progressive(gen, tags):
+def test_storage_upgrade_progressive(gen: str, tags: List[str]):
     PIN = "1234"
 
     def asserts(client):
@@ -148,7 +149,7 @@ def test_storage_upgrade_progressive(gen, tags):
 
 
 @for_all("legacy", legacy_minimum_version=(1, 9, 0))
-def test_upgrade_wipe_code(gen, tag):
+def test_upgrade_wipe_code(gen: str, tag: str):
     PIN = "1234"
     WIPE_CODE = "4321"
 
@@ -193,7 +194,7 @@ def test_upgrade_wipe_code(gen, tag):
 
 
 @for_all("legacy")
-def test_upgrade_reset(gen, tag):
+def test_upgrade_reset(gen: str, tag: str):
     def asserts(client):
         assert not client.features.pin_protection
         assert not client.features.passphrase_protection
@@ -226,7 +227,7 @@ def test_upgrade_reset(gen, tag):
 
 
 @for_all()
-def test_upgrade_reset_skip_backup(gen, tag):
+def test_upgrade_reset_skip_backup(gen: str, tag: str):
     def asserts(client):
         assert not client.features.pin_protection
         assert not client.features.passphrase_protection
@@ -260,7 +261,7 @@ def test_upgrade_reset_skip_backup(gen, tag):
 
 
 @for_all(legacy_minimum_version=(1, 7, 2))
-def test_upgrade_reset_no_backup(gen, tag):
+def test_upgrade_reset_no_backup(gen: str, tag: str):
     def asserts(client):
         assert not client.features.pin_protection
         assert not client.features.passphrase_protection
@@ -295,7 +296,7 @@ def test_upgrade_reset_no_backup(gen, tag):
 
 # 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, tag):
+def test_upgrade_shamir_recovery(gen: str, tag: str):
     with EmulatorWrapper(gen, tag) as emu, BackgroundDeviceHandler(
         emu.client
     ) as device_handler:
@@ -335,7 +336,7 @@ def test_upgrade_shamir_recovery(gen, tag):
 
 
 @for_all(legacy_minimum_version=(1, 8, 4), core_minimum_version=(2, 1, 9))
-def test_upgrade_u2f(gen, tag):
+def test_upgrade_u2f(gen: str, tag: str):
     """Check U2F counter stayed the same after an upgrade."""
     with EmulatorWrapper(gen, tag) as emu:
         debuglink.load_device_by_mnemonic(

tests/upgrade_tests/test_passphrase_consistency.py

@@ -17,6 +17,7 @@
 import pytest
 
 from trezorlib import btc, device, mapping, messages, models, protobuf
+from trezorlib._internal.emulator import Emulator
 from trezorlib.tools import parse_path
 
 from ..emulators import EmulatorWrapper
@@ -40,8 +41,9 @@ mapping.DEFAULT_MAPPING.register(ApplySettingsCompat)
 
 
 @pytest.fixture
-def emulator(gen, tag):
+def emulator(gen: str, tag: str) -> Emulator:
     with EmulatorWrapper(gen, tag) as emu:
+        assert emu.client is not None
         # set up a passphrase-protected device
         device.reset(
             emu.client,
@@ -60,8 +62,9 @@ def emulator(gen, tag):
     core_minimum_version=models.TREZOR_T.minimum_version,
     legacy_minimum_version=models.TREZOR_ONE.minimum_version,
 )
-def test_passphrase_works(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,
@@ -95,10 +98,11 @@ def test_passphrase_works(emulator):
     core_minimum_version=models.TREZOR_T.minimum_version,
     legacy_minimum_version=(1, 9, 0),
 )
-def test_init_device(emulator):
+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,