refactor(core): simplify THP write-related code

diff --git a/core/src/trezor/wire/thp/__init__.py b/core/src/trezor/wire/thp/__init__.py
index fc5e124f9..7e47fab41 100644
--- a/core/src/trezor/wire/thp/__init__.py
+++ b/core/src/trezor/wire/thp/__init__.py
@@ -11,6 +11,7 @@ from ..protocol_common import WireError
 
 if TYPE_CHECKING:
     from enum import IntEnum
+    from typing import Iterable
 
     from trezor.wire import WireInterface
     from typing_extensions import Self
@@ -87,8 +88,11 @@ class SessionState(IntEnum):
 
 
 class PacketHeader:
-    format_str_init = ">BHH"
-    format_str_cont = ">BH"
+    INIT_FORMAT = ">BHH"
+    CONT_FORMAT = ">BH"
+
+    INIT_LENGTH = ustruct.calcsize(INIT_FORMAT)
+    CONT_LENGTH = ustruct.calcsize(CONT_FORMAT)
 
     def __init__(self, ctrl_byte: int, cid: int, length: int) -> None:
         self.ctrl_byte = ctrl_byte
@@ -96,7 +100,7 @@ class PacketHeader:
         self.length = length
 
     def to_bytes(self) -> bytes:
-        return ustruct.pack(self.format_str_init, self.ctrl_byte, self.cid, self.length)
+        return ustruct.pack(self.INIT_FORMAT, self.ctrl_byte, self.cid, self.length)
 
     def pack_to_init_buffer(self, buffer: bytearray, buffer_offset: int = 0) -> None:
         """
@@ -104,7 +108,7 @@ class PacketHeader:
         into the provided buffer.
         """
         ustruct.pack_into(
-            self.format_str_init,
+            self.INIT_FORMAT,
             buffer,
             buffer_offset,
             self.ctrl_byte,
@@ -118,9 +122,38 @@ class PacketHeader:
         into the provided buffer.
         """
         ustruct.pack_into(
-            self.format_str_cont, buffer, buffer_offset, CONTINUATION_PACKET, self.cid
+            self.CONT_FORMAT, buffer, buffer_offset, CONTINUATION_PACKET, self.cid
         )
 
+    def fragment_payload(self, packet_size: int, *items: bytes) -> Iterable[bytes]:
+        """Fragment payload into THP transport packets."""
+        packet = bytearray(packet_size)
+        self.pack_to_init_buffer(packet)
+
+        buf = memoryview(packet)[self.INIT_LENGTH :]
+        buf_offset = 0
+        should_zero_pad = False
+
+        for item in items:
+            item_offset = 0
+            while item_offset < len(item):
+                n = utils.memcpy(buf, buf_offset, item, item_offset)
+                buf_offset += n
+                item_offset += n
+
+                if buf_offset == len(buf):
+                    should_zero_pad = True
+                    yield packet  # packet is full - send to the host
+                    self.pack_to_cont_buffer(packet)
+                    buf = memoryview(packet)[self.CONT_LENGTH :]
+                    buf_offset = 0
+
+        if buf_offset > 0:
+            # send last packet (pad with zeroes if needed)
+            if should_zero_pad:
+                utils.memzero(buf[buf_offset:])
+            yield packet
+
     @classmethod
     def get_error_header(cls, cid: int, length: int) -> Self:
         """
diff --git a/core/src/trezor/wire/thp/channel.py b/core/src/trezor/wire/thp/channel.py
index 873ea58c4..cf1291fe5 100644
--- a/core/src/trezor/wire/thp/channel.py
+++ b/core/src/trezor/wire/thp/channel.py
@@ -27,12 +27,7 @@ from . import alternating_bit_protocol as ABP
 from . import control_byte, crypto, memory_manager
 from .checksum import CHECKSUM_LENGTH
 from .transmission_loop import TransmissionLoop
-from .writer import (
-    CONT_HEADER_LENGTH,
-    INIT_HEADER_LENGTH,
-    MESSAGE_TYPE_LENGTH,
-    write_payload_to_wire_and_add_checksum,
-)
+from .writer import MESSAGE_TYPE_LENGTH
 
 if __debug__:
     from trezor import log
@@ -44,6 +39,7 @@ if TYPE_CHECKING:
     from typing import Any, Awaitable
 
     from trezor.messages import ThpPairingCredential
+    from trezor.wire import WireInterface
 
     from .interface_context import ThpContext
     from .pairing_context import PairingContext
@@ -73,11 +69,11 @@ class Reassembler:
 
             # may raise WireBufferError
             buffer = memory_manager.get_existing_read_buffer(self.cid)
-            self._buffer_packet_data(buffer, packet, CONT_HEADER_LENGTH)
+            self._buffer_packet_data(buffer, packet, PacketHeader.CONT_LENGTH)
         else:
             self.reset()
-            _, _, payload_length = ustruct.unpack(PacketHeader.format_str_init, packet)
-            self.buffer_len = payload_length + INIT_HEADER_LENGTH
+            _, _, payload_length = ustruct.unpack(PacketHeader.INIT_FORMAT, packet)
+            self.buffer_len = payload_length + PacketHeader.INIT_LENGTH
 
             if control_byte.is_ack(ctrl_byte):
                 # don't allocate buffer for ACKs (since they are small)
@@ -113,7 +109,7 @@ class Channel:
 
         # Channel properties
         self.channel_id: bytes = channel_cache.channel_id
-        self.iface = ctx._iface
+        self.ctx: ThpContext = ctx
         if __debug__:
             self._log("channel initialization")
         self.channel_cache: ChannelCache = channel_cache
@@ -134,6 +130,10 @@ class Channel:
         if __debug__:
             self.should_show_pairing_dialog: bool = True
 
+    @property
+    def iface(self) -> WireInterface:
+        return self.ctx._iface
+
     def clear(self) -> None:
         clear_sessions_with_channel_id(self.channel_id)
         memory_manager.release_lock_if_owner(self.get_channel_id_int())
@@ -194,7 +194,7 @@ class Channel:
             raise
 
     def decrypt_buffer(
-        self, message_length: int, offset: int = INIT_HEADER_LENGTH
+        self, message_length: int, offset: int = PacketHeader.INIT_LENGTH
     ) -> None:
         buffer = memory_manager.get_existing_read_buffer(self.get_channel_id_int())
 
@@ -255,7 +255,7 @@ class Channel:
         cid = self.get_channel_id_int()
         msg_size = protobuf.encoded_length(msg)
         payload_size = SESSION_ID_LENGTH + MESSAGE_TYPE_LENGTH + msg_size
-        length = payload_size + CHECKSUM_LENGTH + TAG_LENGTH + INIT_HEADER_LENGTH
+        length = payload_size + CHECKSUM_LENGTH + TAG_LENGTH + PacketHeader.INIT_LENGTH
 
         buffer = memory_manager.get_new_write_buffer(cid, length)
         noise_payload_len = memory_manager.encode_into_buffer(buffer, msg, session_id)
@@ -268,7 +268,7 @@ class Channel:
         msg_data = err_type.to_bytes(1, "big")
         length = len(msg_data) + CHECKSUM_LENGTH
         header = PacketHeader.get_error_header(self.get_channel_id_int(), length)
-        return write_payload_to_wire_and_add_checksum(self.iface, header, msg_data)
+        return self.ctx.write_payload(header, msg_data)
 
     def write_handshake_message(self, ctrl_byte: int, payload: bytes) -> None:
         self._prepare_write()
diff --git a/core/src/trezor/wire/thp/interface_context.py b/core/src/trezor/wire/thp/interface_context.py
index cc54520af..a2ce89f10 100644
--- a/core/src/trezor/wire/thp/interface_context.py
+++ b/core/src/trezor/wire/thp/interface_context.py
@@ -1,5 +1,6 @@
 import ustruct
 from micropython import const
+from trezorcrypto import crc
 from typing import TYPE_CHECKING
 
 from storage.cache_thp import (
@@ -7,7 +8,7 @@ from storage.cache_thp import (
     ChannelCache,
     iter_allocated_channels,
 )
-from trezor import io, loop
+from trezor import io, loop, utils
 
 from ..errors import WireBufferError
 from . import (
@@ -23,20 +24,16 @@ from . import (
 )
 from .channel import Channel
 from .checksum import CHECKSUM_LENGTH
-from .writer import INIT_HEADER_LENGTH, write_payload_to_wire_and_add_checksum
 
 if __debug__:
     from trezor import log
 
 if TYPE_CHECKING:
     from trezorio import WireInterface
+    from typing import Awaitable, Iterable
 
 _CID_REQ_PAYLOAD_LENGTH = const(12)
 
-# Codec_v1 magic constant:
-# "?##" + Failure message type + msg_size + msg_data (code = "Failure_InvalidProtocol")
-CODEC_V1_ERROR_MESSAGE = b"\x3f\x23\x23\x00\x03\x00\x00\x00\x14\x08\x11"
-
 
 class ThpContext:
     """
@@ -96,27 +93,43 @@ class ThpContext:
                 await channel.write_error(ThpErrorType.TRANSPORT_BUSY)
                 continue
 
-    async def write_packet(self, packet: bytes) -> None:
-        assert len(packet) == self._iface.TX_PACKET_LEN
-        while True:
-            await self._write
-            n_written = self._iface.write(packet)
-            if n_written == len(packet):
-                return
-            assert n_written == 0
+    def write_payload(self, header: PacketHeader, payload: bytes) -> Awaitable[None]:
+        checksum = crc.crc32(payload, crc.crc32(header.to_bytes()))
+        checksum_bytes = checksum.to_bytes(CHECKSUM_LENGTH, "big")
+        return self._write_payload_chunks(header, payload, checksum_bytes)
+
+    def _write_payload_chunks(
+        self, header: PacketHeader, *chunks: bytes
+    ) -> Awaitable[None]:
+        fragments = header.fragment_payload(self._iface.TX_PACKET_LEN, *chunks)
+        return self._write_packets(fragments)
+
+    async def _write_packets(self, fragments: Iterable[bytes]) -> None:
+        packet_len = self._iface.TX_PACKET_LEN
+        for packet in fragments:
+            assert len(packet) == packet_len
+
+            n_written = 0
+            while n_written == 0:
+                await self._write
+                n_written = self._iface.write(packet)
+
+            assert n_written == packet_len
 
     async def _handle_codec_v1(self, packet: bytes) -> None:
         # If the received packet is not an initial codec_v1 packet, do not send error message
         if packet[1:3] == b"##":
-            buf = CODEC_V1_ERROR_MESSAGE
-            buf += bytes(self._iface.TX_PACKET_LEN - len(buf))
-            await self.write_packet(buf)
+            response = bytearray(self._iface.TX_PACKET_LEN)
+            # Codec_v1 magic constant:
+            # "?##" + Failure message type + msg_size + msg_data (code = "Failure_InvalidProtocol")
+            utils.memcpy(response, 0, b"?##\x00\x03\x00\x00\x00\x14\x08\x11", 0)
+            await self._write_packets([response])
 
     async def _handle_broadcast(self, packet: bytes) -> None:
         if _get_ctrl_byte(packet) != CHANNEL_ALLOCATION_REQ:
             raise ThpError("Unexpected ctrl_byte in a broadcast channel packet")
 
-        data = packet[: INIT_HEADER_LENGTH + _CID_REQ_PAYLOAD_LENGTH]
+        data = packet[: PacketHeader.INIT_LENGTH + _CID_REQ_PAYLOAD_LENGTH]
         if not checksum.is_valid(data[-CHECKSUM_LENGTH:], data[:-CHECKSUM_LENGTH]):
             raise ThpError("Checksum is not valid")
 
@@ -140,16 +153,14 @@ class ThpContext:
                 channel.get_channel_id_int(),
                 iface=self._iface,
             )
-        await write_payload_to_wire_and_add_checksum(
-            self._iface, response_header, response_data
-        )
+        await self.write_payload(response_header, response_data)
 
     async def _handle_unallocated(self, cid: int, packet: bytes) -> None:
         if control_byte.is_continuation(_get_ctrl_byte(packet)):
             return
         data = (ThpErrorType.UNALLOCATED_CHANNEL).to_bytes(1, "big")
         header = PacketHeader.get_error_header(cid, len(data) + CHECKSUM_LENGTH)
-        await write_payload_to_wire_and_add_checksum(self._iface, header, data)
+        await self.write_payload(header, data)
 
 
 def _get_ctrl_byte(packet: bytes) -> int:
diff --git a/core/src/trezor/wire/thp/received_message_handler.py b/core/src/trezor/wire/thp/received_message_handler.py
index 5b4fb6343..e70fb5b20 100644
--- a/core/src/trezor/wire/thp/received_message_handler.py
+++ b/core/src/trezor/wire/thp/received_message_handler.py
@@ -43,11 +43,7 @@ from . import checksum, control_byte, get_encoded_device_properties, session_man
 from .checksum import CHECKSUM_LENGTH
 from .crypto import PUBKEY_LENGTH, Handshake
 from .session_context import SeedlessSessionContext
-from .writer import (
-    INIT_HEADER_LENGTH,
-    MESSAGE_TYPE_LENGTH,
-    write_payload_to_wire_and_add_checksum,
-)
+from .writer import MESSAGE_TYPE_LENGTH
 
 if TYPE_CHECKING:
     from typing import Awaitable
@@ -84,7 +80,7 @@ async def handle_received_message(
         #         "To show allocation count, create the build with TREZOR_MEMPERF=1"
         #     )
     ctrl_byte, _, payload_length = ustruct.unpack(">BHH", message_buffer)
-    message_length = payload_length + INIT_HEADER_LENGTH
+    message_length = payload_length + PacketHeader.INIT_LENGTH
 
     _check_checksum(message_length, message_buffer)
 
@@ -151,18 +147,18 @@ async def handle_received_message(
         log.debug(__name__, "handle_received_message - end", iface=ctx.iface)
 
 
-def _send_ack(ctx: Channel, ack_bit: int) -> Awaitable[None]:
+def _send_ack(channel: Channel, ack_bit: int) -> Awaitable[None]:
     ctrl_byte = control_byte.add_ack_bit_to_ctrl_byte(ACK_MESSAGE, ack_bit)
-    header = PacketHeader(ctrl_byte, ctx.get_channel_id_int(), CHECKSUM_LENGTH)
+    header = PacketHeader(ctrl_byte, channel.get_channel_id_int(), CHECKSUM_LENGTH)
     if __debug__:
         log.debug(
             __name__,
             "Writing ACK message to a channel with cid: %s, ack_bit: %d",
-            hexlify_if_bytes(ctx.channel_id),
+            hexlify_if_bytes(channel.channel_id),
             ack_bit,
-            iface=ctx.iface,
+            iface=channel.iface,
         )
-    return write_payload_to_wire_and_add_checksum(ctx.iface, header, b"")
+    return channel.ctx.write_payload(header, b"")
 
 
 def _check_checksum(message_length: int, message_buffer: utils.BufferType) -> None:
@@ -255,7 +251,7 @@ def _handle_state_TH1(
     # pass  # TODO buffer is gone :/
 
     host_ephemeral_public_key = bytearray(
-        buffer[INIT_HEADER_LENGTH : message_length - CHECKSUM_LENGTH]
+        buffer[PacketHeader.INIT_LENGTH : message_length - CHECKSUM_LENGTH]
     )
     trezor_ephemeral_public_key, encrypted_trezor_static_public_key, tag = (
         ctx.handshake.handle_th1_crypto(
@@ -306,10 +302,13 @@ def _handle_state_TH2(ctx: Channel, message_length: int, ctrl_byte: int) -> None
     # if buffer is BufferError:
     # pass  # TODO handle
     host_encrypted_static_public_key = buffer[
-        INIT_HEADER_LENGTH : INIT_HEADER_LENGTH + KEY_LENGTH + TAG_LENGTH
+        PacketHeader.INIT_LENGTH : PacketHeader.INIT_LENGTH + KEY_LENGTH + TAG_LENGTH
     ]
     handshake_completion_request_noise_payload = buffer[
-        INIT_HEADER_LENGTH + KEY_LENGTH + TAG_LENGTH : message_length - CHECKSUM_LENGTH
+        PacketHeader.INIT_LENGTH
+        + KEY_LENGTH
+        + TAG_LENGTH : message_length
+        - CHECKSUM_LENGTH
     ]
 
     ctx.handshake.handle_th2_crypto(
@@ -324,7 +323,7 @@ def _handle_state_TH2(ctx: Channel, message_length: int, ctrl_byte: int) -> None
 
     noise_payload = _decode_message(
         buffer[
-            INIT_HEADER_LENGTH
+            PacketHeader.INIT_LENGTH
             + KEY_LENGTH
             + TAG_LENGTH : message_length
             - CHECKSUM_LENGTH
@@ -393,7 +392,7 @@ def _handle_state_ENCRYPTED_TRANSPORT(ctx: Channel, message_length: int) -> None
     # if buffer is BufferError:
     # pass  # TODO handle
     session_id, message_type = ustruct.unpack(
-        ">BH", memoryview(buffer)[INIT_HEADER_LENGTH:]
+        ">BH", memoryview(buffer)[PacketHeader.INIT_LENGTH :]
     )
     if session_id not in ctx.sessions:
 
@@ -415,7 +414,7 @@ def _handle_state_ENCRYPTED_TRANSPORT(ctx: Channel, message_length: int) -> None
         Message(
             message_type,
             buffer[
-                INIT_HEADER_LENGTH
+                PacketHeader.INIT_LENGTH
                 + MESSAGE_TYPE_LENGTH
                 + SESSION_ID_LENGTH : message_length
                 - CHECKSUM_LENGTH
@@ -443,14 +442,14 @@ def _handle_pairing(ctx: Channel, message_length: int) -> None:
     # if buffer is BufferError:
     # pass  # TODO handle
     message_type = ustruct.unpack(
-        ">H", buffer[INIT_HEADER_LENGTH + SESSION_ID_LENGTH :]
+        ">H", buffer[PacketHeader.INIT_LENGTH + SESSION_ID_LENGTH :]
     )[0]
 
     ctx.connection_context.incoming_message.put(
         Message(
             message_type,
             buffer[
-                INIT_HEADER_LENGTH
+                PacketHeader.INIT_LENGTH
                 + MESSAGE_TYPE_LENGTH
                 + SESSION_ID_LENGTH : message_length
                 - CHECKSUM_LENGTH
diff --git a/core/src/trezor/wire/thp/transmission_loop.py b/core/src/trezor/wire/thp/transmission_loop.py
index 1eef89f78..015b9ba99 100644
--- a/core/src/trezor/wire/thp/transmission_loop.py
+++ b/core/src/trezor/wire/thp/transmission_loop.py
@@ -3,8 +3,6 @@ from typing import TYPE_CHECKING
 
 from trezor import loop
 
-from .writer import write_payload_to_wire_and_add_checksum
-
 if TYPE_CHECKING:
     from . import PacketHeader
     from .channel import Channel
@@ -32,9 +30,7 @@ class TransmissionLoop:
         for i in range(max_retransmission_count):
             if i >= MIN_RETRANSMISSION_COUNT:
                 self.min_retransmisson_count_achieved = True
-            await write_payload_to_wire_and_add_checksum(
-                self.channel.iface, self.header, self.transport_payload
-            )
+            await self.channel.ctx.write_payload(self.header, self.transport_payload)
 
             # Do not create wait task for last iteration
             if i == max_retransmission_count - 1:
diff --git a/core/src/trezor/wire/thp/writer.py b/core/src/trezor/wire/thp/writer.py
index 4508ae175..7f13d61a6 100644
--- a/core/src/trezor/wire/thp/writer.py
+++ b/core/src/trezor/wire/thp/writer.py
@@ -1,105 +1,5 @@
 from micropython import const
-from trezorcrypto import crc
-from typing import TYPE_CHECKING
 
-from trezor import io, loop, utils
-
-from . import PacketHeader
-
-INIT_HEADER_LENGTH = const(5)
-CONT_HEADER_LENGTH = const(3)
 CHECKSUM_LENGTH = const(4)
 MAX_PAYLOAD_LEN = const(60000)
 MESSAGE_TYPE_LENGTH = const(2)
-
-if TYPE_CHECKING:
-    from trezorio import WireInterface
-    from typing import Awaitable, Sequence
-
-if __debug__:
-    from trezor import log
-
-
-def write_payload_to_wire_and_add_checksum(
-    iface: WireInterface, header: PacketHeader, transport_payload: bytes
-) -> Awaitable[None]:
-    header_checksum: int = crc.crc32(header.to_bytes())
-    checksum: bytes = crc.crc32(transport_payload, header_checksum).to_bytes(
-        CHECKSUM_LENGTH, "big"
-    )
-    data = (transport_payload, checksum)
-    return write_payloads_to_wire(iface, header, data)
-
-
-async def write_payloads_to_wire(
-    iface: WireInterface, header: PacketHeader, data: Sequence[bytes]
-) -> None:
-    n_of_data = len(data)
-    total_length = sum(len(item) for item in data)
-
-    current_data_idx = 0
-    current_data_offset = 0
-
-    packet = bytearray(iface.TX_PACKET_LEN)
-    header.pack_to_init_buffer(packet)
-    packet_offset: int = INIT_HEADER_LENGTH
-    packet_number = 0
-    nwritten = 0
-    while nwritten < total_length:
-        if packet_number >= 1 and nwritten >= total_length - iface.TX_PACKET_LEN:
-            # zero the last packet bytes and add the cont header
-            packet[:] = bytearray(iface.TX_PACKET_LEN)
-            header.pack_to_cont_buffer(packet)
-        elif packet_number == 1:
-            # add the cont header when preparing the second packet (if it is not the last)
-            header.pack_to_cont_buffer(packet)
-
-        while True:
-            n = utils.memcpy(
-                packet, packet_offset, data[current_data_idx], current_data_offset
-            )
-            packet_offset += n
-            current_data_offset += n
-            nwritten += n
-
-            if packet_offset < iface.TX_PACKET_LEN:
-                current_data_idx += 1
-                current_data_offset = 0
-                if current_data_idx >= n_of_data:
-                    break
-            elif packet_offset == iface.TX_PACKET_LEN:
-                break
-            else:
-                raise Exception  # Should not happen
-        packet_number += 1
-        packet_offset = CONT_HEADER_LENGTH
-
-        # write packet to wire (in-lined)
-        if __debug__:
-            log.debug(
-                __name__,
-                "write_packet_to_wire: %s",
-                utils.hexlify_if_bytes(packet),
-                iface=iface,
-            )
-        while True:
-            await loop.wait(iface.iface_num() | io.POLL_WRITE)
-            written_by_iface = iface.write(packet)
-            if written_by_iface == len(packet):
-                break
-            assert written_by_iface == 0
-
-
-async def write_packet_to_wire(iface: WireInterface, packet: bytes) -> None:
-    while True:
-        await loop.wait(iface.iface_num() | io.POLL_WRITE)
-        if __debug__:
-            log.debug(
-                __name__,
-                "write_packet_to_wire: %s",
-                utils.hexlify_if_bytes(packet),
-                iface=iface,
-            )
-        n_written = iface.write(packet)
-        if n_written == len(packet):
-            return
diff --git a/core/tests/test_trezor.wire.thp.writer.py b/core/tests/test_trezor.wire.thp.writer.py
index 14abba31c..bbbdea4e1 100644
--- a/core/tests/test_trezor.wire.thp.writer.py
+++ b/core/tests/test_trezor.wire.thp.writer.py
@@ -6,7 +6,8 @@ from typing import Any, Awaitable
 if utils.USE_THP:
     import thp_common
     from mock_wire_interface import MockHID
-    from trezor.wire.thp import ENCRYPTED, PacketHeader, writer
+    from trezor.wire.thp import ENCRYPTED, PacketHeader
+    from trezor.wire.thp.interface_context import ThpContext
 
 
 @unittest.skipUnless(utils.USE_THP, "only needed for THP")
@@ -76,33 +77,23 @@ class TestTrezorHostProtocolWriter(unittest.TestCase):
 
     def setUp(self):
         self.interface = MockHID()
-
-    def test_write_empty_packet(self):
-        self.await_until_result(writer.write_packet_to_wire(self.interface, b""))
-
-        print(self.interface.data[0])
-        self.assertEqual(len(self.interface.data), 1)
-        self.assertEqual(self.interface.data[0], b"")
+        self.ctx = ThpContext(self.interface)
 
     def test_write_empty_payload(self):
         header = PacketHeader(ENCRYPTED, 4660, 4)
-        await_result(writer.write_payloads_to_wire(self.interface, header, (b"",)))
+        await_result(self.ctx._write_payload_chunks(header, b""))
         self.assertEqual(len(self.interface.data), 0)
 
     def test_write_short_payload(self):
         header = PacketHeader(ENCRYPTED, 4660, 5)
         data = b"\x07"
-        self.await_until_result(
-            writer.write_payloads_to_wire(self.interface, header, (data,))
-        )
+        self.await_until_result(self.ctx._write_payload_chunks(header, data))
         self.assertEqual(hexlify(self.interface.data[0]), self.short_payload_expected)
 
     def test_write_longer_payload(self):
         data = bytearray(range(256))
         header = PacketHeader(ENCRYPTED, 4660, 256)
-        self.await_until_result(
-            writer.write_payloads_to_wire(self.interface, header, (data,))
-        )
+        self.await_until_result(self.ctx._write_payload_chunks(header, data))
 
         for i in range(len(self.longer_payload_expected)):
             self.assertEqual(
@@ -112,11 +103,9 @@ class TestTrezorHostProtocolWriter(unittest.TestCase):
     def test_write_eight_longer_payloads(self):
         data = bytearray(range(256))
         header = PacketHeader(ENCRYPTED, 4660, 2048)
-        self.await_until_result(
-            writer.write_payloads_to_wire(
-                self.interface, header, (data, data, data, data, data, data, data, data)
-            )
-        )
+        chunks = [data] * 8
+        self.await_until_result(self.ctx._write_payload_chunks(header, *chunks))
+
         for i in range(len(self.eight_longer_payloads_expected)):
             self.assertEqual(
                 hexlify(self.interface.data[i]), self.eight_longer_payloads_expected[i]
@@ -124,9 +113,7 @@ class TestTrezorHostProtocolWriter(unittest.TestCase):
 
     def test_write_empty_payload_with_checksum(self):
         header = PacketHeader(ENCRYPTED, 4660, 4)
-        self.await_until_result(
-            writer.write_payload_to_wire_and_add_checksum(self.interface, header, b"")
-        )
+        self.await_until_result(self.ctx.write_payload(header, b""))
 
         self.assertEqual(
             hexlify(self.interface.data[0]), self.empty_payload_with_checksum_expected
@@ -135,9 +122,7 @@ class TestTrezorHostProtocolWriter(unittest.TestCase):
     def test_write_longer_payload_with_checksum(self):
         data = bytearray(range(256))
         header = PacketHeader(ENCRYPTED, 4660, 256)
-        self.await_until_result(
-            writer.write_payload_to_wire_and_add_checksum(self.interface, header, data)
-        )
+        self.await_until_result(self.ctx.write_payload(header, data))
 
         for i in range(len(self.longer_payload_with_checksum_expected)):
             self.assertEqual(