tests: simplify change_wipe_code test with PIN queuing

tests/device_tests/test_msg_change_wipe_code_t1.py

@@ -16,7 +16,9 @@
 
 import pytest
 
-from trezorlib import messages
+from trezorlib import device, exceptions, messages
+
+PinType = messages.PinMatrixRequestType
 
 PIN4 = "1234"
 WIPE_CODE4 = "4321"
@@ -27,77 +29,46 @@ pytestmark = pytest.mark.skip_t2
 
 def _set_wipe_code(client, wipe_code):
     # Set/change wipe code.
-    ret = client.call_raw(messages.ChangeWipeCode())
-    assert isinstance(ret, messages.ButtonRequest)
-
-    # Confirm intent to set/change wipe code.
-    client.debug.press_yes()
-    ret = client.call_raw(messages.ButtonAck())
-
-    if client.features.pin_protection:
-        # Send current PIN.
-        assert isinstance(ret, messages.PinMatrixRequest)
-        pin_encoded = client.debug.read_pin_encoded()
-        ret = client.call_raw(messages.PinMatrixAck(pin=pin_encoded))
-
-    # Send the new wipe code for the first time.
-    assert isinstance(ret, messages.PinMatrixRequest)
-    wipe_code_encoded = client.debug.encode_pin(wipe_code)
-    ret = client.call_raw(messages.PinMatrixAck(pin=wipe_code_encoded))
-
-    # Send the new wipe code for the second time.
-    assert isinstance(ret, messages.PinMatrixRequest)
-    wipe_code_encoded = client.debug.encode_pin(wipe_code)
-    ret = client.call_raw(messages.PinMatrixAck(pin=wipe_code_encoded))
-
-    # Now we're done.
-    assert isinstance(ret, messages.Success)
-
-
-def _remove_wipe_code(client):
-    # Remove wipe code
-    ret = client.call_raw(messages.ChangeWipeCode(remove=True))
-    assert isinstance(ret, messages.ButtonRequest)
-
-    # Confirm intent to remove wipe code.
-    client.debug.press_yes()
-    ret = client.call_raw(messages.ButtonAck())
-
-    # Send current PIN.
-    assert isinstance(ret, messages.PinMatrixRequest)
-    pin_encoded = client.debug.read_pin_encoded()
-    ret = client.call_raw(messages.PinMatrixAck(pin=pin_encoded))
-
-    # Now we're done.
-    assert isinstance(ret, messages.Success)
+    with client:
+        if client.features.pin_protection:
+            pin, _ = client.debug.read_pin()
+            pins = [pin, wipe_code, wipe_code]
+            pin_matrices = [
+                messages.PinMatrixRequest(type=PinType.Current),
+                messages.PinMatrixRequest(type=PinType.WipeCodeFirst),
+                messages.PinMatrixRequest(type=PinType.WipeCodeSecond),
+            ]
+        else:
+            pins = [wipe_code, wipe_code]
+            pin_matrices = [
+                messages.PinMatrixRequest(type=PinType.WipeCodeFirst),
+                messages.PinMatrixRequest(type=PinType.WipeCodeSecond),
+            ]
+
+        client.set_pin(pins)
+        client.set_expected_responses(
+            [messages.ButtonRequest()]
+            + pin_matrices
+            + [messages.Success(), messages.Features()]
+        )
+        device.change_wipe_code(client)
+
+
+def _change_pin(client, old_pin, new_pin):
+    assert client.features.pin_protection is True
+    with client:
+        client.set_pin([old_pin, new_pin, new_pin])
+        try:
+            return device.change_pin(client)
+        except exceptions.TrezorFailure as f:
+            return f.failure
 
 
 def _check_wipe_code(client, wipe_code):
-    # Try to change the PIN to the current wipe code value. The operation should fail.
-    ret = client.call_raw(messages.ChangePin())
-    assert isinstance(ret, messages.ButtonRequest)
-
-    # Confirm intent to change PIN.
-    client.debug.press_yes()
-    ret = client.call_raw(messages.ButtonAck())
-
-    # Send current PIN.
-    assert isinstance(ret, messages.PinMatrixRequest)
-    pin_encoded = client.debug.read_pin_encoded()
-    ret = client.call_raw(messages.PinMatrixAck(pin=pin_encoded))
-
-    # Send the new wipe code for the first time.
-    assert isinstance(ret, messages.PinMatrixRequest)
-    wipe_code_encoded = client.debug.encode_pin(wipe_code)
-    ret = client.call_raw(messages.PinMatrixAck(pin=wipe_code_encoded))
-
-    # Send the new wipe code for the second time.
-    assert isinstance(ret, messages.PinMatrixRequest)
-    wipe_code_encoded = client.debug.encode_pin(wipe_code)
-    ret = client.call_raw(messages.PinMatrixAck(pin=wipe_code_encoded))
-
-    # Expect failure.
-    assert isinstance(ret, messages.Failure)
+    """Check that wipe code is set by changing the PIN to it."""
+    old_pin, _ = client.debug.read_pin()
+    f = _change_pin(client, old_pin, wipe_code)
+    assert isinstance(f, messages.Failure)
 
 
 @pytest.mark.setup_client(pin=PIN4)
@@ -122,11 +93,11 @@ def test_set_remove_wipe_code(client):
     client.init_device()
     assert client.features.wipe_code_protection is True
 
-    # Check that the PIN is correct.
+    # Check that the wipe code is correct.
     _check_wipe_code(client, WIPE_CODE6)
 
     # Test remove wipe code.
-    _remove_wipe_code(client)
+    device.change_wipe_code(client, remove=True)
 
     # Check that there's no wipe code protection now.
     client.init_device()
@@ -138,26 +109,18 @@ def test_set_wipe_code_mismatch(client):
     assert client.features.wipe_code_protection is False
 
     # Let's set a new wipe code.
-    ret = client.call_raw(messages.ChangeWipeCode())
-    assert isinstance(ret, messages.ButtonRequest)
-
-    # Confirm intent to set wipe code.
-    client.debug.press_yes()
-    ret = client.call_raw(messages.ButtonAck())
-
-    # Send the new wipe code for the first time.
-    assert isinstance(ret, messages.PinMatrixRequest)
-    wipe_code_encoded = client.debug.encode_pin(WIPE_CODE4)
-    ret = client.call_raw(messages.PinMatrixAck(pin=wipe_code_encoded))
-
-    # Send the new wipe code for the second time, but different.
-    assert isinstance(ret, messages.PinMatrixRequest)
-    wipe_code_encoded = client.debug.encode_pin(WIPE_CODE6)
-    ret = client.call_raw(messages.PinMatrixAck(pin=wipe_code_encoded))
-
-    # The operation should fail, because the wipe codes are different.
-    assert isinstance(ret, messages.Failure)
-    assert ret.code == messages.FailureType.WipeCodeMismatch
+    with client:
+        client.set_pin([WIPE_CODE4, WIPE_CODE6])
+        client.set_expected_responses(
+            [
+                messages.ButtonRequest(),
+                messages.PinMatrixRequest(type=PinType.WipeCodeFirst),
+                messages.PinMatrixRequest(type=PinType.WipeCodeSecond),
+                messages.Failure(code=messages.FailureType.WipeCodeMismatch),
+            ]
+        )
+        with pytest.raises(exceptions.TrezorFailure):
+            device.change_wipe_code(client)
 
     # Check that there is no wipe code protection.
     client.init_device()
@@ -170,26 +133,18 @@ def test_set_wipe_code_to_pin(client):
     assert client.features.wipe_code_protection is None
 
     # Let's try setting the wipe code to the curent PIN value.
-    ret = client.call_raw(messages.ChangeWipeCode())
-    assert isinstance(ret, messages.ButtonRequest)
-
-    # Confirm intent to set wipe code.
-    client.debug.press_yes()
-    ret = client.call_raw(messages.ButtonAck())
-
-    # Send current PIN.
-    assert isinstance(ret, messages.PinMatrixRequest)
-    pin_encoded = client.debug.read_pin_encoded()
-    ret = client.call_raw(messages.PinMatrixAck(pin=pin_encoded))
-
-    # Send the new wipe code.
-    assert isinstance(ret, messages.PinMatrixRequest)
-    pin_encoded = client.debug.read_pin_encoded()
-    ret = client.call_raw(messages.PinMatrixAck(pin=pin_encoded))
-
-    # The operation should fail, because the wipe code must be different from the PIN.
-    assert isinstance(ret, messages.Failure)
-    assert ret.code == messages.FailureType.ProcessError
+    with client:
+        client.set_pin([PIN4, PIN4])
+        client.set_expected_responses(
+            [
+                messages.ButtonRequest(),
+                messages.PinMatrixRequest(type=PinType.Current),
+                messages.PinMatrixRequest(type=PinType.WipeCodeFirst),
+                messages.Failure(code=messages.FailureType.ProcessError),
+            ]
+        )
+        with pytest.raises(exceptions.TrezorFailure):
+            device.change_wipe_code(client)
 
     # Check that there is no wipe code protection.
     client.init_device()
@@ -201,26 +156,18 @@ def test_set_pin_to_wipe_code(client):
     _set_wipe_code(client, WIPE_CODE4)
 
     # Try to set the PIN to the current wipe code value.
-    ret = client.call_raw(messages.ChangePin())
-    assert isinstance(ret, messages.ButtonRequest)
-
-    # Confirm intent to set PIN.
-    client.debug.press_yes()
-    ret = client.call_raw(messages.ButtonAck())
-
-    # Send the new PIN for the first time.
-    assert isinstance(ret, messages.PinMatrixRequest)
-    pin_encoded = client.debug.encode_pin(WIPE_CODE4)
-    ret = client.call_raw(messages.PinMatrixAck(pin=pin_encoded))
-
-    # Send the new PIN for the second time.
-    assert isinstance(ret, messages.PinMatrixRequest)
-    pin_encoded = client.debug.encode_pin(WIPE_CODE4)
-    ret = client.call_raw(messages.PinMatrixAck(pin=pin_encoded))
-
-    # The operation should fail, because the PIN must be different from the wipe code.
-    assert isinstance(ret, messages.Failure)
-    assert ret.code == messages.FailureType.ProcessError
+    with client:
+        client.set_pin([WIPE_CODE4, WIPE_CODE4])
+        client.set_expected_responses(
+            [
+                messages.ButtonRequest(),
+                messages.PinMatrixRequest(type=PinType.NewFirst),
+                messages.PinMatrixRequest(type=PinType.NewSecond),
+                messages.Failure(code=messages.FailureType.ProcessError),
+            ]
+        )
+        with pytest.raises(exceptions.TrezorFailure):
+            device.change_pin(client)
 
     # Check that there is no PIN protection.
     client.init_device()