trezor-agent/trezor_agent/tests/test_util.py

import io

import pytest

from .. import util


def test_bytes2num():
    assert util.bytes2num(b'\x12\x34') == 0x1234


def test_num2bytes():
    assert util.num2bytes(0x1234, size=2) == b'\x12\x34'


def test_pack():
    assert util.pack('BHL', 1, 2, 3) == b'\x01\x00\x02\x00\x00\x00\x03'


def test_frames():
    msgs = [b'aaa', b'bb', b'c' * 0x12340]
    f = util.frame(*msgs)
    assert f == b'\x00\x01\x23\x45' + b''.join(msgs)
    assert util.read_frame(io.BytesIO(f)) == b''.join(msgs)


class FakeSocket(object):
    def __init__(self):
        self.buf = io.BytesIO()

    def sendall(self, data):
        self.buf.write(data)

    def recv(self, size):
        return self.buf.read(size)


def test_send_recv():
    s = FakeSocket()
    util.send(s, b'123')
    util.send(s, b'*')
    assert s.buf.getvalue() == b'123*'

    s.buf.seek(0)
    assert util.recv(s, 2) == b'12'
    assert util.recv(s, 2) == b'3*'

    pytest.raises(EOFError, util.recv, s, 1)