diff --git a/tests/persistence_tests/test_safety_checks.py b/tests/persistence_tests/test_safety_checks.py index 1cbf7d7551..c2552f04b3 100644 --- a/tests/persistence_tests/test_safety_checks.py +++ b/tests/persistence_tests/test_safety_checks.py @@ -20,16 +20,18 @@ from ..upgrade_tests import core_only def test_safety_checks_level_after_reboot( core_emulator: Emulator, set_level: SafetyCheckLevel, after_level: SafetyCheckLevel ): - device.wipe(core_emulator.client) + device.wipe(core_emulator.client.get_management_session()) + core_emulator.client = core_emulator.client.get_new_client() debuglink.load_device( - core_emulator.client, + core_emulator.client.get_management_session(), mnemonic=MNEMONIC12, pin="", passphrase_protection=False, label="SAFETYLEVEL", ) - device.apply_settings(core_emulator.client, safety_checks=set_level) + device.apply_settings(core_emulator.client.get_session(), safety_checks=set_level) + core_emulator.client.refresh_features() assert core_emulator.client.features.safety_checks == set_level core_emulator.restart() diff --git a/tests/persistence_tests/test_shamir_persistence.py b/tests/persistence_tests/test_shamir_persistence.py index e24f16eeb6..1524bd5203 100644 --- a/tests/persistence_tests/test_shamir_persistence.py +++ b/tests/persistence_tests/test_shamir_persistence.py @@ -16,7 +16,8 @@ import pytest -from trezorlib import device +from trezorlib import device, messages +from trezorlib.client import ProtocolVersion from trezorlib.debuglink import DebugLink, LayoutType from trezorlib.messages import RecoveryStatus @@ -45,7 +46,7 @@ def test_abort(core_emulator: Emulator): assert features.recovery_status == RecoveryStatus.Nothing - device_handler.run(device.recover, pin_protection=False) + device_handler.run_with_session(device.recover, pin_protection=False) recovery.confirm_recovery(debug) layout = debug.read_layout() @@ -82,7 +83,7 @@ def test_recovery_single_reset(core_emulator: Emulator): assert features.initialized is False assert features.recovery_status == RecoveryStatus.Nothing - device_handler.run(device.recover, pin_protection=False) + device_handler.run_with_session(device.recover, pin_protection=False) recovery.confirm_recovery(debug) @@ -129,7 +130,7 @@ def test_recovery_on_old_wallet(core_emulator: Emulator): assert features.recovery_status == RecoveryStatus.Nothing # enter recovery mode - device_handler.run(device.recover, pin_protection=False) + device_handler.run_with_session(device.recover, pin_protection=False) recovery.confirm_recovery(debug) @@ -157,7 +158,8 @@ def test_recovery_on_old_wallet(core_emulator: Emulator): layout = debug.read_layout() # while keyboard is open, hit the device with Initialize/GetFeatures - device_handler.client.init_device() + if device_handler.client.protocol_version == ProtocolVersion.PROTOCOL_V1: + device_handler.client.get_management_session().call(messages.Initialize()) device_handler.client.refresh_features() # try entering remaining 19 words @@ -207,7 +209,7 @@ def test_recovery_multiple_resets(core_emulator: Emulator): assert features.recovery_status == RecoveryStatus.Nothing # start device and recovery - device_handler.run(device.recover, pin_protection=False) + device_handler.run_with_session(device.recover, pin_protection=False) recovery.confirm_recovery(debug) diff --git a/tests/persistence_tests/test_wipe_code.py b/tests/persistence_tests/test_wipe_code.py index 2497a708f6..cb06eeb2cd 100644 --- a/tests/persistence_tests/test_wipe_code.py +++ b/tests/persistence_tests/test_wipe_code.py @@ -11,46 +11,55 @@ WIPE_CODE = "9876" def setup_device_legacy(client: Client, pin: str, wipe_code: str) -> None: - device.wipe(client) + device.wipe(client.get_management_session()) + client = client.get_new_client() debuglink.load_device( - client, MNEMONIC12, pin, passphrase_protection=False, label="WIPECODE" + client.get_management_session(), + MNEMONIC12, + pin, + passphrase_protection=False, + label="WIPECODE", ) with client: client.use_pin_sequence([PIN, WIPE_CODE, WIPE_CODE]) - device.change_wipe_code(client) + device.change_wipe_code(client.get_management_session()) def setup_device_core(client: Client, pin: str, wipe_code: str) -> None: - device.wipe(client) + device.wipe(client.get_management_session()) + client = client.get_new_client() debuglink.load_device( - client, MNEMONIC12, pin, passphrase_protection=False, label="WIPECODE" + client.get_management_session(), + MNEMONIC12, + pin, + passphrase_protection=False, + label="WIPECODE", ) with client: client.use_pin_sequence([pin, wipe_code, wipe_code]) - device.change_wipe_code(client) + device.change_wipe_code(client.get_management_session()) @core_only def test_wipe_code_activate_core(core_emulator: Emulator): # set up device setup_device_core(core_emulator.client, PIN, WIPE_CODE) - - core_emulator.client.init_device() + session = core_emulator.client.get_session() device_id = core_emulator.client.features.device_id # Initiate Change pin process - ret = core_emulator.client.call_raw(messages.ChangePin(remove=False)) + ret = session.call_raw(messages.ChangePin(remove=False)) assert isinstance(ret, messages.ButtonRequest) assert ret.name == "change_pin" core_emulator.client.debug.press_yes() - ret = core_emulator.client.call_raw(messages.ButtonAck()) + ret = session.call_raw(messages.ButtonAck()) # Enter the wipe code instead of the current PIN expected = message_filters.ButtonRequest(code=messages.ButtonRequestType.PinEntry) assert expected.match(ret) - core_emulator.client._raw_write(messages.ButtonAck()) + session._write(messages.ButtonAck()) core_emulator.client.debug.input(WIPE_CODE) # preserving screenshots even after it dies and starts again @@ -75,25 +84,26 @@ def test_wipe_code_activate_legacy(): # set up device setup_device_legacy(emu.client, PIN, WIPE_CODE) - emu.client.init_device() + session = emu.client.get_session() device_id = emu.client.features.device_id # Initiate Change pin process - ret = emu.client.call_raw(messages.ChangePin(remove=False)) + ret = session.call_raw(messages.ChangePin(remove=False)) assert isinstance(ret, messages.ButtonRequest) emu.client.debug.press_yes() - ret = emu.client.call_raw(messages.ButtonAck()) + ret = session.call_raw(messages.ButtonAck()) # Enter the wipe code instead of the current PIN assert isinstance(ret, messages.PinMatrixRequest) wipe_code_encoded = emu.client.debug.encode_pin(WIPE_CODE) - emu.client._raw_write(messages.PinMatrixAck(pin=wipe_code_encoded)) + session._write(messages.PinMatrixAck(pin=wipe_code_encoded)) # wait 30 seconds for emulator to shut down # this will raise a TimeoutError if the emulator doesn't die. emu.wait(30) emu.start() + emu.client.refresh_features() assert emu.client.features.initialized is False assert emu.client.features.pin_protection is False assert emu.client.features.wipe_code_protection is False