diff --git a/tests/device_tests/bitcoin/test_getaddress_show.py b/tests/device_tests/bitcoin/test_getaddress_show.py index bcb685db1d..464c9cc70e 100644 --- a/tests/device_tests/bitcoin/test_getaddress_show.py +++ b/tests/device_tests/bitcoin/test_getaddress_show.py @@ -59,13 +59,13 @@ def test_show_t1( ): def input_flow_t1(): yield - session.debug.press_no() + session.client.debug.press_no() yield - session.debug.press_yes() + session.client.debug.press_yes() - with session: + with session.client as client: # This is the only place where even T1 is using input flow - session.set_input_flow(input_flow_t1) + client.set_input_flow(input_flow_t1) assert ( btc.get_address( session, diff --git a/tests/device_tests/bitcoin/test_getpublickey.py b/tests/device_tests/bitcoin/test_getpublickey.py index 7fcc1a595e..e013e6f71c 100644 --- a/tests/device_tests/bitcoin/test_getpublickey.py +++ b/tests/device_tests/bitcoin/test_getpublickey.py @@ -160,13 +160,13 @@ def test_get_public_node_show_legacy( with client: # test XPUB display flow (without showing QR code) - res = btc.get_public_node(client, path, coin_name=coin_name, show_display=True) + res = btc.get_public_node(session, path, coin_name=coin_name, show_display=True) assert res.xpub == xpub assert bip32.serialize(res.node, xpub_magic) == xpub # test XPUB QR code display using the input flow above client.set_input_flow(input_flow) - res = btc.get_public_node(client, path, coin_name=coin_name, show_display=True) + res = btc.get_public_node(session, path, coin_name=coin_name, show_display=True) assert res.xpub == xpub assert bip32.serialize(res.node, xpub_magic) == xpub diff --git a/tests/device_tests/reset_recovery/test_recovery_bip39_t1.py b/tests/device_tests/reset_recovery/test_recovery_bip39_t1.py index 7ddc634b8d..5f6b17242d 100644 --- a/tests/device_tests/reset_recovery/test_recovery_bip39_t1.py +++ b/tests/device_tests/reset_recovery/test_recovery_bip39_t1.py @@ -77,9 +77,9 @@ def test_pin_passphrase(session: Session): assert fakes == 12 assert mnemonic == [None] * 12 - raise Exception("TEST IS USING INIT MESSAGE - TODO CHANGE") # Mnemonic is the same - session.init_device() + session.init_session() + session.client.refresh_features() assert debug.state().mnemonic_secret == MNEMONIC12.encode() assert session.features.pin_protection is True @@ -131,10 +131,9 @@ def test_nopin_nopassphrase(session: Session): assert fakes == 12 assert mnemonic == [None] * 12 - raise Exception("TEST IS USING INIT MESSAGE - TODO CHANGE") - # Mnemonic is the same - # session.init_device() + session.init_session() + session.client.refresh_features() assert debug.state().mnemonic_secret == MNEMONIC12.encode() assert session.features.pin_protection is False diff --git a/tests/device_tests/reset_recovery/test_reset_bip39_t1.py b/tests/device_tests/reset_recovery/test_reset_bip39_t1.py index b482a5af69..ef4cc264b8 100644 --- a/tests/device_tests/reset_recovery/test_reset_bip39_t1.py +++ b/tests/device_tests/reset_recovery/test_reset_bip39_t1.py @@ -54,8 +54,8 @@ def reset_device(session: Session, strength: int): mnemonic = [] for _ in range(strength // 32 * 3): assert isinstance(ret, messages.ButtonRequest) - mnemonic.append(session.debug.read_reset_word()) - session.debug.press_yes() + mnemonic.append(session.client.debug.read_reset_word()) + session.client.debug.press_yes() session.call_raw(messages.ButtonAck()) mnemonic = " ".join(mnemonic) @@ -66,7 +66,7 @@ def reset_device(session: Session, strength: int): mnemonic = [] for _ in range(strength // 32 * 3): assert isinstance(ret, messages.ButtonRequest) - mnemonic.append(session.debug.read_reset_word()) + mnemonic.append(session.client.debug.read_reset_word()) debug.press_yes() resp = session.call_raw(messages.ButtonAck()) @@ -78,7 +78,7 @@ def reset_device(session: Session, strength: int): assert mnemonic == expected_mnemonic # Check if device is properly initialized - resp = session.call_raw(messages.GetFeatures()) + resp = session.call_raw(messages.Initialize()) assert resp.initialized is True assert resp.backup_availability == messages.BackupAvailability.NotAvailable assert resp.pin_protection is False @@ -170,7 +170,7 @@ def test_reset_device_256_pin(session: Session): assert mnemonic == expected_mnemonic # Check if device is properly initialized - resp = session.call_raw(messages.GetFeatures()) + resp = session.call_raw(messages.Initialize()) assert resp.initialized is True assert resp.backup_availability == messages.BackupAvailability.NotAvailable assert resp.pin_protection is True diff --git a/tests/device_tests/test_msg_change_wipe_code_t1.py b/tests/device_tests/test_msg_change_wipe_code_t1.py index 033f11be60..5a4a17d60c 100644 --- a/tests/device_tests/test_msg_change_wipe_code_t1.py +++ b/tests/device_tests/test_msg_change_wipe_code_t1.py @@ -19,7 +19,6 @@ import pytest from trezorlib import device, exceptions, messages from trezorlib.client import MAX_PIN_LENGTH from trezorlib.debuglink import SessionDebugWrapper as Session -from trezorlib.debuglink import TrezorClientDebugLink as Client from trezorlib.tools import parse_path PinType = messages.PinMatrixRequestType @@ -33,10 +32,10 @@ WIPE_CODE_TOO_LONG = WIPE_CODE_MAX + "1" pytestmark = pytest.mark.models("legacy") -def _set_wipe_code(client: Client, pin, wipe_code): +def _set_wipe_code(session: Session, pin, wipe_code): # Set/change wipe code. - with client: - if client.features.pin_protection: + with session.client as client, session: + if session.features.pin_protection: pins = [pin, wipe_code, wipe_code] pin_matrices = [ messages.PinMatrixRequest(type=PinType.Current), @@ -51,63 +50,58 @@ def _set_wipe_code(client: Client, pin, wipe_code): ] client.use_pin_sequence(pins) - client.set_expected_responses( - [messages.ButtonRequest()] - + pin_matrices - + [messages.Success, messages.Features] + session.set_expected_responses( + [messages.ButtonRequest()] + pin_matrices + [messages.Success] ) - device.change_wipe_code(client) + device.change_wipe_code(session) -def _change_pin(client: Client, old_pin, new_pin): - assert client.features.pin_protection is True - with client: +def _change_pin(session: Session, old_pin, new_pin): + assert session.features.pin_protection is True + with session.client as client: client.use_pin_sequence([old_pin, new_pin, new_pin]) try: - return device.change_pin(client) + return device.change_pin(session) except exceptions.TrezorFailure as f: return f.failure -def _check_wipe_code(client: Client, pin, wipe_code): +def _check_wipe_code(session: Session, pin, wipe_code): """Check that wipe code is set by changing the PIN to it.""" - f = _change_pin(client, pin, wipe_code) + f = _change_pin(session, pin, wipe_code) assert isinstance(f, messages.Failure) @pytest.mark.setup_client(pin=PIN4) -def test_set_remove_wipe_code(client: Client): +def test_set_remove_wipe_code(session: Session): # Check that wipe code protection status is not revealed in locked state. - assert client.features.wipe_code_protection is None + assert session.features.wipe_code_protection is None # Test set wipe code. - _set_wipe_code(client, PIN4, WIPE_CODE_MAX) + _set_wipe_code(session, PIN4, WIPE_CODE_MAX) # Check that there's wipe code protection now. - client.init_device() - assert client.features.wipe_code_protection is True + assert session.features.wipe_code_protection is True # Check that the wipe code is correct. - _check_wipe_code(client, PIN4, WIPE_CODE_MAX) + _check_wipe_code(session, PIN4, WIPE_CODE_MAX) # Test change wipe code. - _set_wipe_code(client, PIN4, WIPE_CODE6) + _set_wipe_code(session, PIN4, WIPE_CODE6) # Check that there's still wipe code protection now. - client.init_device() - assert client.features.wipe_code_protection is True + assert session.features.wipe_code_protection is True # Check that the wipe code is correct. - _check_wipe_code(client, PIN4, WIPE_CODE6) + _check_wipe_code(session, PIN4, WIPE_CODE6) # Test remove wipe code. - with client: + with session.client as client: client.use_pin_sequence([PIN4]) - device.change_wipe_code(client, remove=True) + device.change_wipe_code(session, remove=True) # Check that there's no wipe code protection now. - client.init_device() - assert client.features.wipe_code_protection is False + assert session.features.wipe_code_protection is False def test_set_wipe_code_mismatch(session: Session): @@ -136,14 +130,14 @@ def test_set_wipe_code_mismatch(session: Session): @pytest.mark.setup_client(pin=PIN4) -def test_set_wipe_code_to_pin(client: Client): +def test_set_wipe_code_to_pin(session: Session): # Check that wipe code protection status is not revealed in locked state. - assert client.features.wipe_code_protection is None + assert session.features.wipe_code_protection is None # Let's try setting the wipe code to the curent PIN value. - with client: + with session.client as client, session: client.use_pin_sequence([PIN4, PIN4]) - client.set_expected_responses( + session.set_expected_responses( [ messages.ButtonRequest(), messages.PinMatrixRequest(type=PinType.Current), @@ -152,21 +146,22 @@ def test_set_wipe_code_to_pin(client: Client): ] ) with pytest.raises(exceptions.TrezorFailure): - device.change_wipe_code(client) + device.change_wipe_code(session) # Check that there is no wipe code protection. - client.init_device() + client.refresh_features() assert client.features.wipe_code_protection is False -def test_set_pin_to_wipe_code(client: Client): +def test_set_pin_to_wipe_code(session: Session): # Set wipe code. - _set_wipe_code(client, None, WIPE_CODE4) + session.refresh_features() + _set_wipe_code(session, None, WIPE_CODE4) # Try to set the PIN to the current wipe code value. - with client: + with session.client as client, session: client.use_pin_sequence([WIPE_CODE4, WIPE_CODE4]) - client.set_expected_responses( + session.set_expected_responses( [ messages.ButtonRequest(), messages.PinMatrixRequest(type=PinType.NewFirst), @@ -175,17 +170,18 @@ def test_set_pin_to_wipe_code(client: Client): ] ) with pytest.raises(exceptions.TrezorFailure): - device.change_pin(client) + device.change_pin(session) # Check that there is no PIN protection. - client.init_device() - assert client.features.pin_protection is False - resp = client.call_raw(messages.GetAddress(address_n=parse_path("m/44'/0'/0'/0/0"))) + assert session.features.pin_protection is False + resp = session.call_raw( + messages.GetAddress(address_n=parse_path("m/44'/0'/0'/0/0")) + ) assert isinstance(resp, messages.Address) @pytest.mark.parametrize("invalid_wipe_code", ("1204", "", WIPE_CODE_TOO_LONG)) -def test_set_wipe_code_invalid(session: Session, invalid_wipe_code): +def test_set_wipe_code_invalid(session: Session, invalid_wipe_code: str): # Let's set the wipe code ret = session.call_raw(messages.ChangeWipeCode()) assert isinstance(ret, messages.ButtonRequest) @@ -203,6 +199,6 @@ def test_set_wipe_code_invalid(session: Session, invalid_wipe_code): assert isinstance(ret, messages.Failure) # Check that there's still no wipe code protection. - session.refresh_features() + session.init_session() # why no response on GetFeatures? session.ensure_unlocked() assert session.features.wipe_code_protection is False diff --git a/tests/device_tests/test_msg_changepin_t1.py b/tests/device_tests/test_msg_changepin_t1.py index b60820eee7..7ca59d2979 100644 --- a/tests/device_tests/test_msg_changepin_t1.py +++ b/tests/device_tests/test_msg_changepin_t1.py @@ -61,7 +61,6 @@ def test_set_pin(session: Session): messages.PinMatrixRequest, messages.PinMatrixRequest, messages.Success, - messages.Features, ] ) device.change_pin(session) @@ -88,7 +87,6 @@ def test_change_pin(session: Session): messages.PinMatrixRequest, messages.PinMatrixRequest, messages.Success, - messages.Features, ] ) device.change_pin(session) @@ -113,7 +111,6 @@ def test_remove_pin(session: Session): messages.ButtonRequest(code=messages.ButtonRequestType.ProtectCall), messages.PinMatrixRequest, messages.Success, - messages.Features, ] ) device.change_pin(session, remove=True) @@ -199,7 +196,7 @@ def test_set_invalid(session: Session, invalid_pin): # Check that there's still no PIN protection now # TODO change session.init_device() - session.refresh_features() + session.init_session() assert session.features.pin_protection is False _check_no_pin(session) diff --git a/tests/device_tests/test_protection_levels.py b/tests/device_tests/test_protection_levels.py index c73975ffcb..5d5d2a02e6 100644 --- a/tests/device_tests/test_protection_levels.py +++ b/tests/device_tests/test_protection_levels.py @@ -133,7 +133,6 @@ def test_change_pin_t1(session: Session): _pin_request(session), _pin_request(session), messages.Success, - messages.Features, ] ) device.change_pin(session) diff --git a/tests/device_tests/test_session_id_and_passphrase.py b/tests/device_tests/test_session_id_and_passphrase.py index 153915cad4..defbb0340b 100644 --- a/tests/device_tests/test_session_id_and_passphrase.py +++ b/tests/device_tests/test_session_id_and_passphrase.py @@ -324,16 +324,18 @@ def test_passphrase_always_on_device(client: Client): @pytest.mark.models("legacy") @pytest.mark.setup_client(passphrase="") -def test_passphrase_on_device_not_possible_on_t1(client: Client): +def test_passphrase_on_device_not_possible_on_t1(session: Session): # This setting makes no sense on T1. - response = client.call_raw(messages.ApplySettings(passphrase_always_on_device=True)) + response = session.call_raw( + messages.ApplySettings(passphrase_always_on_device=True) + ) assert isinstance(response, messages.Failure) assert response.code == FailureType.DataError # T1 should not accept on_device request - response = client.call_raw(XPUB_REQUEST) + response = session.call_raw(XPUB_REQUEST) assert isinstance(response, messages.PassphraseRequest) - response = client.call_raw(messages.PassphraseAck(on_device=True)) + response = session.call_raw(messages.PassphraseAck(on_device=True)) assert isinstance(response, messages.Failure) assert response.code == FailureType.DataError