tests: add wipe code activation test

tests/device_tests/test_msg_change_wipe_code_t2.py

@@ -219,42 +219,3 @@ def test_set_pin_to_wipe_code(client):
         )
         client.set_input_flow(_input_flow_set_pin(client.debug, WIPE_CODE4))
         device.change_pin(client)
-
-
-@pytest.mark.setup_client(pin=PIN4)
-def test_wipe_code_activate(client):
-    import time
-
-    device_id = client.features.device_id
-
-    # Set wipe code.
-    with client:
-        client.set_expected_responses(
-            [messages.ButtonRequest()] * 5 + [messages.Success(), messages.Features()]
-        )
-        client.set_input_flow(_input_flow_set_wipe_code(client.debug, PIN4, WIPE_CODE4))
-
-        device.change_wipe_code(client)
-
-    # Try to change the PIN.
-    ret = client.call_raw(messages.ChangePin(remove=False))
-
-    # Confirm change PIN.
-    assert isinstance(ret, messages.ButtonRequest)
-    client.debug.press_yes()
-    ret = client.call_raw(messages.ButtonAck())
-
-    # Enter the wipe code instead of the current PIN
-    assert ret == messages.ButtonRequest(code=messages.ButtonRequestType.Other)
-    client.debug.input(WIPE_CODE4)
-    client._raw_write(messages.ButtonAck())
-
-    # Allow the device to display wipe code popup and restart.
-    time.sleep(7)
-
-    # Check that the device has been wiped.
-    client.init_device()
-    assert client.features.initialized is False
-    assert client.features.pin_protection is False
-    assert client.features.wipe_code_protection is False
-    assert client.features.device_id != device_id

tests/persistence_tests/test_wipe_code.py

@@ -0,0 +1,108 @@
+from trezorlib import debuglink, device, messages
+
+from ..common import MNEMONIC12
+from ..emulators import EmulatorWrapper
+from ..upgrade_tests import core_only, legacy_only
+
+PIN = "1234"
+WIPE_CODE = "9876"
+
+
+def setup_device_legacy(client, pin, wipe_code):
+    device.wipe(client)
+    debuglink.load_device(
+        client, MNEMONIC12, pin, passphrase_protection=False, label="WIPECODE"
+    )
+
+    with client:
+        client.set_pin([PIN, WIPE_CODE, WIPE_CODE])
+        device.change_wipe_code(client)
+
+
+def setup_device_core(client, pin, wipe_code):
+    device.wipe(client)
+    debuglink.load_device(
+        client, MNEMONIC12, pin, passphrase_protection=False, label="WIPECODE"
+    )
+
+    def input_flow():
+        yield  # do you want to set/change the wipe_code?
+        client.debug.press_yes()
+        if pin is not None:
+            yield  # enter current pin
+            client.debug.input(pin)
+        yield  # enter new wipe code
+        client.debug.input(wipe_code)
+        yield  # enter new wipe code again
+        client.debug.input(wipe_code)
+        yield  # success
+        client.debug.press_yes()
+
+    with client:
+        client.set_expected_responses(
+            [messages.ButtonRequest()] * 5 + [messages.Success(), messages.Features()]
+        )
+        client.set_input_flow(input_flow)
+        device.change_wipe_code(client)
+
+
+@core_only
+def test_wipe_code_activate_core():
+    with EmulatorWrapper("core") as emu:
+        # set up device
+        setup_device_core(emu.client, PIN, WIPE_CODE)
+
+        emu.client.init_device()
+        device_id = emu.client.features.device_id
+
+        # Initiate Change pin process
+        ret = emu.client.call_raw(messages.ChangePin(remove=False))
+        assert isinstance(ret, messages.ButtonRequest)
+        emu.client.debug.press_yes()
+        ret = emu.client.call_raw(messages.ButtonAck())
+
+        # Enter the wipe code instead of the current PIN
+        assert ret == messages.ButtonRequest(code=messages.ButtonRequestType.Other)
+        emu.client._raw_write(messages.ButtonAck())
+        emu.client.debug.input(WIPE_CODE)
+
+        # wait 30 seconds for emulator to shut down
+        # this will raise a TimeoutError if the emulator doesn't die.
+        emu.wait(30)
+
+        emu.start()
+        assert emu.client.features.initialized is False
+        assert emu.client.features.pin_protection is False
+        assert emu.client.features.wipe_code_protection is False
+        assert emu.client.features.device_id != device_id
+
+
+@legacy_only
+def test_wipe_code_activate_legacy():
+    with EmulatorWrapper("legacy") as emu:
+        # set up device
+        setup_device_legacy(emu.client, PIN, WIPE_CODE)
+
+        emu.client.init_device()
+        device_id = emu.client.features.device_id
+
+        # Initiate Change pin process
+        ret = emu.client.call_raw(messages.ChangePin(remove=False))
+        assert isinstance(ret, messages.ButtonRequest)
+        emu.client.debug.press_yes()
+        ret = emu.client.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))
+
+        # wait 30 seconds for emulator to shut down
+        # this will raise a TimeoutError if the emulator doesn't die.
+        emu.wait(30)
+
+        emu.start()
+        assert emu.client.features.initialized is False
+        assert emu.client.features.pin_protection is False
+        assert emu.client.features.wipe_code_protection is False
+        assert emu.client.features.device_id != device_id