test: update persistence tests

[no changelog]

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()

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)
 

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