test(core): fix all THP tests, unify usage of mock interface

[no changelog]

core/tests/mock_wire_interface.py

@@ -0,0 +1,17 @@
+from trezor.loop import wait
+
+
+class MockHID:
+    def __init__(self, num):
+        self.num = num
+        self.data = []
+
+    def iface_num(self):
+        return self.num
+
+    def write(self, msg):
+        self.data.append(bytearray(msg))
+        return len(msg)
+
+    def wait_object(self, mode):
+        return wait(mode | self.num)

core/tests/test_apps.bitcoin.approver.py

@@ -11,6 +11,7 @@ from trezor.messages import (
     TxInput,
     TxOutput,
 )
+from trezor.wire import context
 
 from apps.bitcoin.authorization import FEE_RATE_DECIMALS, CoinJoinAuthorization
 from apps.bitcoin.sign_tx.approvers import CoinJoinApprover
@@ -18,17 +19,22 @@ from apps.bitcoin.sign_tx.bitcoin import Bitcoin
 from apps.bitcoin.sign_tx.tx_info import TxInfo
 from apps.common import coins
 
-if not utils.USE_THP:
+if utils.USE_THP:
+    from thp_common import prepare_context, suppres_debug_log
+else:
     import storage.cache_codec
 
 class TestApprover(unittest.TestCase):
-    if not utils.USE_THP:
+    if utils.USE_THP:
 
         def __init__(self):
-            # Context is needed to test decorators and handleInitialize
-            # It allows access to codec cache from different parts of the code
-            from trezor.wire import context
+            suppres_debug_log()
+            prepare_context()
+            super().__init__()
 
+    else:
+
+        def __init__(self):
             context.CURRENT_CONTEXT = context.CodecContext(None, bytearray(64))
             super().__init__()
 

core/tests/test_apps.bitcoin.authorization.py

@@ -3,24 +3,30 @@ from common import *  # isort:skip
 import storage.cache
 from trezor.enums import InputScriptType
 from trezor.messages import AuthorizeCoinJoin, GetOwnershipProof, SignTx
+from trezor.wire import context
 
 from apps.bitcoin.authorization import CoinJoinAuthorization
 from apps.common import coins
 
 _ROUND_ID_LEN = 32
 
-if not utils.USE_THP:
+if utils.USE_THP:
+    from thp_common import prepare_context, suppres_debug_log
+else:
     import storage.cache_codec
 
 
 class TestAuthorization(unittest.TestCase):
-    if not utils.USE_THP:
+    if utils.USE_THP:
 
         def __init__(self):
-            # Context is needed to test decorators and handleInitialize
-            # It allows access to codec cache from different parts of the code
-            from trezor.wire import context
+            suppres_debug_log()
+            prepare_context()
+            super().__init__()
 
+    else:
+
+        def __init__(self):
             context.CURRENT_CONTEXT = context.CodecContext(None, bytearray(64))
             super().__init__()
 

core/tests/test_apps.bitcoin.keychain.py

@@ -3,21 +3,31 @@ from common import *  # isort:skip
 from storage import cache_common
 from trezor import wire
 from trezor.crypto import bip39
+from trezor.wire import context
 
 from apps.bitcoin.keychain import _get_coin_by_name, _get_keychain_for_coin
 
-if not utils.USE_THP:
+if utils.USE_THP:
+    from thp_common import prepare_context, suppres_debug_log
+else:
     from storage import cache_codec
 
 
 class TestBitcoinKeychain(unittest.TestCase):
-    if not utils.USE_THP:
+    if utils.USE_THP:
 
         def __init__(self):
-            # Context is needed to test decorators and handleInitialize
-            # It allows access to codec cache from different parts of the code
-            from trezor.wire import context
+            suppres_debug_log()
+            prepare_context()
+            super().__init__()
 
+        def setUp(self):
+            seed = bip39.seed(" ".join(["all"] * 12), "")
+            context.cache_set(cache_common.APP_COMMON_SEED, seed)
+
+    else:
+
+        def __init__(self):
             context.CURRENT_CONTEXT = context.CodecContext(None, bytearray(64))
             super().__init__()
 

core/tests/test_apps.common.keychain.py

@@ -5,30 +5,39 @@ from storage import cache, cache_common
 from trezor import utils, wire
 from trezor.crypto import bip39
 from trezor.enums import SafetyCheckLevel
+from trezor.wire import context
 
 from apps.common import safety_checks
 from apps.common.keychain import Keychain, LRUCache, get_keychain, with_slip44_keychain
 from apps.common.paths import PATTERN_SEP5, PathSchema
 
+if utils.USE_THP:
+    from thp_common import prepare_context, suppres_debug_log
 if not utils.USE_THP:
     from storage import cache_codec
 
 
 class TestKeychain(unittest.TestCase):
 
-    if not utils.USE_THP:
+    if utils.USE_THP:
 
         def __init__(self):
-            # Context is needed to test decorators and handleInitialize
-            # It allows access to codec cache from different parts of the code
-            from trezor.wire import context
+            suppres_debug_log()
+            prepare_context()
+            super().__init__()
 
+    else:
+
+        def __init__(self):
             context.CURRENT_CONTEXT = context.CodecContext(None, bytearray(64))
             super().__init__()
 
         def setUp(self):
             cache_codec.start_session()
 
+    def cache_set(self, key: int, value: bytes) -> None:
+        context.cache_set(key, value)
+
     def tearDown(self):
         cache.clear_all()
 
@@ -85,7 +94,7 @@ class TestKeychain(unittest.TestCase):
 
     def test_get_keychain(self):
         seed = bip39.seed(" ".join(["all"] * 12), "")
-        cache_codec.get_active_session().set(cache_common.APP_COMMON_SEED, seed)
+        self.cache_set(cache_common.APP_COMMON_SEED, seed)
 
         schema = PathSchema.parse("m/44'/1'", 0)
         keychain = await_result(get_keychain("secp256k1", [schema]))
@@ -99,7 +108,7 @@ class TestKeychain(unittest.TestCase):
 
     def test_with_slip44(self):
         seed = bip39.seed(" ".join(["all"] * 12), "")
-        cache_codec.get_active_session().set(cache_common.APP_COMMON_SEED, seed)
+        self.cache_set(cache_common.APP_COMMON_SEED, seed)
 
         slip44_id = 42
         valid_path = [H_(44), H_(slip44_id), H_(0)]

core/tests/test_apps.ethereum.keychain.py

@@ -5,11 +5,14 @@ import unittest
 from storage import cache_common
 from trezor import wire
 from trezor.crypto import bip39
+from trezor.wire import context
 
 from apps.common.keychain import get_keychain
 from apps.common.paths import HARDENED
 
-if not utils.USE_THP:
+if utils.USE_THP:
+    from thp_common import prepare_context, suppres_debug_log
+else:
     from storage import cache_codec
 
 
@@ -74,13 +77,20 @@ class TestEthereumKeychain(unittest.TestCase):
                 keychain.derive,
                 addr,
             )
-    if not utils.USE_THP:
+
+    if utils.USE_THP:
+        def __init__(self):
+            suppres_debug_log()
+            prepare_context()
+            super().__init__()
+
+        def setUp(self):
+            seed = bip39.seed(" ".join(["all"] * 12), "")
+            context.cache_set(cache_common.APP_COMMON_SEED, seed)
+
+    else:
 
         def __init__(self):
-            # Context is needed to test decorators and handleInitialize
-            # It allows access to codec cache from different parts of the code
-            from trezor.wire import context
-
             context.CURRENT_CONTEXT = context.CodecContext(None, bytearray(64))
             super().__init__()
 

core/tests/test_trezor.wire.codec_v1.py

@@ -2,28 +2,11 @@ from common import *  # isort:skip
 
 import ustruct
 
+from mock_wire_interface import MockHID
 from trezor import io
-from trezor.loop import wait
 from trezor.utils import chunks
 from trezor.wire import codec_v1
 
-
-class MockHID:
-    def __init__(self, num):
-        self.num = num
-        self.data = []
-
-    def iface_num(self):
-        return self.num
-
-    def write(self, msg):
-        self.data.append(bytearray(msg))
-        return len(msg)
-
-    def wait_object(self, mode):
-        return wait(mode | self.num)
-
-
 MESSAGE_TYPE = 0x4242
 
 HEADER_PAYLOAD_LENGTH = codec_v1._REP_LEN - 3 - ustruct.calcsize(">HL")

core/tests/test_trezor.wire.thp.py

@@ -1,9 +1,9 @@
 from common import *  # isort:skip
+from mock_wire_interface import MockHID
 from storage import cache_thp
 from trezor import config, io, log, protobuf
 from trezor.crypto.curve import curve25519
 from trezor.enums import MessageType
-from trezor.loop import wait
 from trezor.wire.errors import UnexpectedMessage
 from trezor.wire.protocol_common import Message
 
@@ -34,20 +34,6 @@ if utils.USE_THP:
 
 # Disable log.debug for the test
 log.debug = lambda name, msg, *args: None
-class MockHID:
-    def __init__(self, num):
-        self.num = num
-        self.data = []
-
-    def iface_num(self):
-        return self.num
-
-    def write(self, msg):
-        self.data.append(bytearray(msg))
-        return len(msg)
-
-    def wait_object(self, mode):
-        return wait(mode | self.num)
 
 
 def dummy_decode_iface(cached_iface: bytes):

core/tests/test_trezor.wire.thp.writer.py

@@ -3,24 +3,10 @@ from common import *  # isort:skip
 from typing import Any, Awaitable
 
 if utils.USE_THP:
+    from mock_wire_interface import MockHID
     from trezor.wire.thp import writer
     from trezor.wire.thp.thp_messages import ENCRYPTED_TRANSPORT, PacketHeader
 
-    class MockHID:
-        def __init__(self, num):
-            self.num = num
-            self.data = []
-
-        def iface_num(self):
-            return self.num
-
-        def write(self, msg):
-            self.data.append(bytearray(msg))
-            return len(msg)
-
-        def wait_object(self, mode):
-            return wait(mode | self.num)
-
 
 if __debug__:
     # Disable log.debug for the test

core/tests/test_trezor.wire.thp_deprecated.py

@@ -2,9 +2,9 @@ from common import *  # isort:skip
 import ustruct
 from typing import TYPE_CHECKING
 
+from mock_wire_interface import MockHID
 from storage.cache_thp import BROADCAST_CHANNEL_ID
 from trezor import io, log
-from trezor.loop import wait
 from trezor.utils import chunks
 from trezor.wire.protocol_common import Message
 
@@ -22,21 +22,6 @@ if __debug__:
 if TYPE_CHECKING:
     from trezorio import WireInterface
 
-class MockHID:
-    def __init__(self, num):
-        self.num = num
-        self.data = []
-
-    def iface_num(self):
-        return self.num
-
-    def write(self, msg):
-        self.data.append(bytearray(msg))
-        return len(msg)
-
-    def wait_object(self, mode):
-        return wait(mode | self.num)
-
 
 MESSAGE_TYPE = 0x4242
 MESSAGE_TYPE_BYTES = b"\x42\x42"

core/tests/thp_common.py

@@ -0,0 +1,29 @@
+from trezor import utils
+
+if utils.USE_THP:
+    from mock_wire_interface import MockHID
+    from storage import cache_thp
+    from trezor.wire import context
+    from trezor.wire.thp import interface_manager
+    from trezor.wire.thp.channel import Channel
+    from trezor.wire.thp.interface_manager import _MOCK_INTERFACE_HID
+    from trezor.wire.thp.session_context import SessionContext
+
+    def dummy_decode_iface(cached_iface: bytes):
+        return MockHID(0xDEADBEEF)
+
+    def prepare_context() -> None:
+        interface_manager.decode_iface = dummy_decode_iface
+        channel_cache = cache_thp.get_new_channel(_MOCK_INTERFACE_HID)
+        channel = Channel(channel_cache)
+        session_cache = cache_thp.get_new_session(channel_cache)
+        session_ctx = SessionContext(channel, session_cache)
+        context.CURRENT_CONTEXT = session_ctx
+
+
+if __debug__:
+    # Disable log.debug for the test
+    def suppres_debug_log() -> None:
+        from trezor import log
+
+        log.debug = lambda name, msg, *args: None