Telethon/client/tests/transport/intermediate_test.py
2024-03-16 19:20:01 +01:00

62 lines
1.6 KiB
Python

from typing import Tuple
from pytest import raises
from telethon._impl.mtproto import Intermediate
class Output(bytearray):
__slots__ = ()
def __call__(self, data: bytes | bytearray | memoryview) -> None:
self += data
def setup_pack(n: int) -> Tuple[Intermediate, bytes, Output]:
input = bytes(x & 0xFF for x in range(n))
return Intermediate(), input, Output()
def test_pack_empty() -> None:
transport, input, output = setup_pack(0)
transport.pack(input, output)
assert output == b"\xee\xee\xee\xee\0\0\0\0"
def test_pack_non_padded() -> None:
transport, input, output = setup_pack(7)
with raises(AssertionError):
transport.pack(input, output)
def test_pack_normal() -> None:
transport, input, output = setup_pack(128)
transport.pack(input, output)
assert output[:8] == b"\xee\xee\xee\xee\x80\0\0\0"
assert output[8:] == input
def test_unpack_small() -> None:
transport = Intermediate()
input = b"\x01"
output = bytearray()
with raises(ValueError) as e:
transport.unpack(input, output)
e.match("missing bytes")
def test_unpack_normal() -> None:
transport, input, packed = setup_pack(128)
unpacked = bytearray()
transport.pack(input, packed)
transport.unpack(packed[4:], unpacked)
assert input == unpacked
def test_unpack_two_at_once() -> None:
transport, input, packed = setup_pack(128)
unpacked = bytearray()
transport.pack(input, packed)
n = transport.unpack(packed[4:] + packed[4:], unpacked)
assert input == unpacked
assert n == len(packed[4:])