From 5ad8bfd1b89ca53543be1e52fc65e222311129b6 Mon Sep 17 00:00:00 2001
From: M1nd3r <petrsedlacek.km@seznam.cz>
Date: Thu, 2 Jan 2025 17:00:34 +0100
Subject: [PATCH] fix device tests [no changelog]

---
 tests/device_tests/bitcoin/test_multisig.py   |  4 ++--
 .../reset_recovery/test_reset_bip39_t2.py     | 23 +++++++++++++------
 .../reset_recovery/test_reset_slip39_basic.py |  2 +-
 3 files changed, 19 insertions(+), 10 deletions(-)

diff --git a/tests/device_tests/bitcoin/test_multisig.py b/tests/device_tests/bitcoin/test_multisig.py
index 77c528d21a..c6a548709a 100644
--- a/tests/device_tests/bitcoin/test_multisig.py
+++ b/tests/device_tests/bitcoin/test_multisig.py
@@ -214,10 +214,10 @@ def test_pubkeys_order(session: Session):
     )
 
     pubkey_internal = btc.get_public_node(
-        client, parse_path("m/45h/0/0/0"), coin_name="Bitcoin"
+        session, parse_path("m/45h/0/0/0"), coin_name="Bitcoin"
     ).node.public_key
     pubkey_external = btc.get_public_node(
-        client, parse_path("m/45h/1/0/0"), coin_name="Bitcoin"
+        session, parse_path("m/45h/1/0/0"), coin_name="Bitcoin"
     ).node.public_key
 
     # This assertion implies that script pubkey of multisig_sorted_1, multisig_sorted_2 and multisig_unsorted_1 are the same
diff --git a/tests/device_tests/reset_recovery/test_reset_bip39_t2.py b/tests/device_tests/reset_recovery/test_reset_bip39_t2.py
index f985f745ad..61b1a04d12 100644
--- a/tests/device_tests/reset_recovery/test_reset_bip39_t2.py
+++ b/tests/device_tests/reset_recovery/test_reset_bip39_t2.py
@@ -19,6 +19,7 @@ from mnemonic import Mnemonic
 
 from trezorlib import device, messages
 from trezorlib.btc import get_public_node
+from trezorlib.client import ProtocolVersion
 from trezorlib.debuglink import LayoutType
 from trezorlib.debuglink import SessionDebugWrapper as Session
 from trezorlib.exceptions import TrezorFailure
@@ -116,6 +117,7 @@ def test_reset_device_pin(session: Session):
     assert resp.passphrase_protection is True
 
 
+@pytest.mark.setup_client(uninitialized=True)
 @pytest.mark.uninitialized_session
 def test_reset_entropy_check(session: Session):
     strength = 128  # 12 words
@@ -143,19 +145,26 @@ def test_reset_entropy_check(session: Session):
     assert IF.mnemonic == expected_mnemonic
 
     # Check that the device is properly initialized.
-    resp = session.call_raw(messages.Initialize())
-    assert resp.initialized is True
-    assert resp.backup_availability == messages.BackupAvailability.NotAvailable
-    assert resp.pin_protection is False
-    assert resp.passphrase_protection is False
-    assert resp.backup_type is messages.BackupType.Bip39
+    if client.protocol_version is ProtocolVersion.PROTOCOL_V1:
+        features = session.call_raw(messages.Initialize())
+    else:
+        session.refresh_features()
+        features = session.features
+    features
+    assert features.initialized is True
+    assert features.backup_availability == messages.BackupAvailability.NotAvailable
+    assert features.pin_protection is False
+    assert features.passphrase_protection is False
+    assert features.backup_type is messages.BackupType.Bip39
 
     # Check that the XPUBs are the same as those from the entropy check.
+    session = session.client.get_session()
     for path, xpub in path_xpubs:
-        res = get_public_node(client, path)
+        res = get_public_node(session, path)
         assert res.xpub == xpub
 
 
+@pytest.mark.setup_client(uninitialized=True)
 @pytest.mark.uninitialized_session
 def test_reset_failed_check(session: Session):
     debug = session.client.debug
diff --git a/tests/device_tests/reset_recovery/test_reset_slip39_basic.py b/tests/device_tests/reset_recovery/test_reset_slip39_basic.py
index 3e65b80007..84e922e071 100644
--- a/tests/device_tests/reset_recovery/test_reset_slip39_basic.py
+++ b/tests/device_tests/reset_recovery/test_reset_slip39_basic.py
@@ -116,7 +116,7 @@ def test_reset_entropy_check(session: Session):
 
     # Check that the XPUBs are the same as those from the entropy check.
     for path, xpub in path_xpubs:
-        res = get_public_node(client, path)
+        res = get_public_node(session, path)
         assert res.xpub == xpub