From 32ed782fef20bcb7b7c2b53b31962db1fdfb7709 Mon Sep 17 00:00:00 2001 From: Jakob Scheid Date: Thu, 2 Apr 2026 12:34:21 +0200 Subject: [PATCH] Add tests for jebp_utils --- src/jeb_utils/jebp_utils.py | 9 +++-- tests/unit/jeb_utils/test_pack_fields.py | 20 ++++++++++ tests/unit/jeb_utils/test_readmsg.py | 44 ++++++++++++++++++++++ tests/unit/jeb_utils/test_sendmsg.py | 29 ++++++++++++++ tests/unit/jeb_utils/test_unpack_fields.py | 21 +++++++++++ 5 files changed, 119 insertions(+), 4 deletions(-) create mode 100644 tests/unit/jeb_utils/test_pack_fields.py create mode 100644 tests/unit/jeb_utils/test_readmsg.py create mode 100644 tests/unit/jeb_utils/test_sendmsg.py create mode 100644 tests/unit/jeb_utils/test_unpack_fields.py diff --git a/src/jeb_utils/jebp_utils.py b/src/jeb_utils/jebp_utils.py index 3d1faee..271407e 100644 --- a/src/jeb_utils/jebp_utils.py +++ b/src/jeb_utils/jebp_utils.py @@ -90,7 +90,7 @@ async def readmsg(reader: asyncio.StreamReader, aesgcm: AESGCM = None, aesnonce: return aesgcm.decrypt(aesnonce, m, None) return m -def pack_fields(*fields: bytes) -> bytes: +def pack_fields(fields: list[bytes]) -> bytes: ''' Packs the fields into a compact bytestring. @@ -100,11 +100,12 @@ def pack_fields(*fields: bytes) -> bytes: :return: The result :rtype: bytes ''' + result = b'' for field in fields: - field_length = utils.int_to_bytes(len(field)) - field_length_length = utils.int_to_bytes(len(field_length)) - result += field_length_length + field_length + field + payload_length = utils.int_to_bytes(len(field)) + payload_length_length = len(payload_length).to_bytes(1) + result += payload_length_length + payload_length + field return result def unpack_fields(raw: bytes) -> Sequence[bytes]: diff --git a/tests/unit/jeb_utils/test_pack_fields.py b/tests/unit/jeb_utils/test_pack_fields.py new file mode 100644 index 0000000..6bb0446 --- /dev/null +++ b/tests/unit/jeb_utils/test_pack_fields.py @@ -0,0 +1,20 @@ +from src.jeb_utils.jebp_utils import pack_fields +import pytest + +@pytest.mark.parametrize('fields,expected', [ + ([], b''), + ([b''], b'\x01\x00'), + ([b'', b''], b'\x01\x00\x01\x00'), + ([b'', b'a'], b'\x01\x00\x01\x01a'), + ([b'', b'ab'], b'\x01\x00\x01\x02ab'), + ([b'a', b'ab'], b'\x01\x01a\x01\x02ab'), + ([b'a', b'ab', b'', b'Hello, World!'], b'\x01\x01a\x01\x02ab\x01\x00\x01\x0dHello, World!'), + ([b'a' * 255], b'\x01\xff' + b'a' * 255), + ([b'a' * 255, b'a' * 255], (b'\x01\xff' + b'a' * 255) * 2), + ([b'a' * 256], b'\x02\x01\x00' + b'a' * 256), + ([b'a' * 256, b'a' * 256], (b'\x02\x01\x00' + b'a' * 256) * 2), + ([b'a' * 257], b'\x02\x01\x01' + b'a' * 257), + ([b'a' * 257, b'abc'], b'\x02\x01\x01' + b'a' * 257 + b'\x01\x03abc'), +]) +def test_pack_fields(fields, expected): + assert pack_fields(fields) == expected \ No newline at end of file diff --git a/tests/unit/jeb_utils/test_readmsg.py b/tests/unit/jeb_utils/test_readmsg.py new file mode 100644 index 0000000..0e6cd6d --- /dev/null +++ b/tests/unit/jeb_utils/test_readmsg.py @@ -0,0 +1,44 @@ +from src.jeb_utils.jebp_utils import readmsg, _MSG_START_BYTE, MessageFormatError +import pytest + +class AsyncTestStreamReader: + def __init__(self, content: bytes) -> None: + self.content = content + + async def readexactly(self, n: int) -> bytes: + if len(self.content) < n: + raise EOFError + result = self.content[:n] + self.content = self.content[n:] + return result + + +@pytest.mark.asyncio +@pytest.mark.parametrize('expected,reader,start_byte', [ + (b'', AsyncTestStreamReader(b'\x01\x01\x00'), b'\x01'), + (b'', AsyncTestStreamReader(b'\x02\x01\x00'), b'\x02'), + (b'a', AsyncTestStreamReader(b'\x02\x01\x01a'), b'\x02'), + (b'', AsyncTestStreamReader(b'\x01\x00'), b'\x01'), + (b'a' * 255, AsyncTestStreamReader(b'\x01\x01\xff' + b'a' * 255), b'\x01'), + (b'a' * 256, AsyncTestStreamReader(b'\x01\x02\x01\x00' + b'a' * 256), b'\x01'), + (b'a' * 256, AsyncTestStreamReader(b'\x01\x02\x01\x00' + b'a' * 257), b'\x01'), + (b'a' * 257, AsyncTestStreamReader(b'\x01\x02\x01\x01' + b'a' * 257), b'\x01'), +]) +async def test_readmsg(expected, reader, start_byte): + assert await readmsg(reader, start_byte = start_byte) == expected + + +@pytest.mark.asyncio +@pytest.mark.parametrize('expected_exception,reader,start_byte', [ + (EOFError, AsyncTestStreamReader(b'\x01\x01\x01'), b'\x01'), + (EOFError, AsyncTestStreamReader(b'\x01\x02\x01a'), b'\x01'), + (EOFError, AsyncTestStreamReader(b'\x01'), b'\x01'), + (EOFError, AsyncTestStreamReader(b'\x01\x01'), b'\x01'), + (EOFError, AsyncTestStreamReader(b'\x01\x02\x00'), b'\x01'), + (MessageFormatError, AsyncTestStreamReader(b'\x01\x02\x01a'), b'\x02'), + (MessageFormatError, AsyncTestStreamReader(b'\x01\x01\x01'), b'\x02'), + (MessageFormatError, AsyncTestStreamReader(b'\x01'), b'\x02'), +]) +async def test_readmsg_exceptions(expected_exception, reader, start_byte): + with pytest.raises(expected_exception): + await readmsg(reader, start_byte = start_byte) \ No newline at end of file diff --git a/tests/unit/jeb_utils/test_sendmsg.py b/tests/unit/jeb_utils/test_sendmsg.py new file mode 100644 index 0000000..ec7881d --- /dev/null +++ b/tests/unit/jeb_utils/test_sendmsg.py @@ -0,0 +1,29 @@ +from src.jeb_utils.jebp_utils import sendmsg +import pytest + +class AsyncTestStreamWriter: + def __init__(self): + self.content = b'' + + def write(self, data): + self.content += data + + async def drain(self): + pass + + +@pytest.mark.asyncio +@pytest.mark.parametrize('expected,message,send_headers,_content_length,start_byte', [ + (b'\x01\x01\x00', b'', True, None, b'\x01'), + (b'\x02\x01\x00', b'', True, None, b'\x02'), + (b'\x01\x01\x01a', b'a', True, None, b'\x01'), + (b'\x01a', b'a', False, None, b'\x01'), + (b'\x01a', b'a', False, 42, b'\x01'), + (b'\x01\x01\x2aa', b'a', True, 42, b'\x01'), + (b'\x01\x0209a', b'a', True, 12345, b'\x01'), + (b'\x01a', b'a', False, 42, b'\x01'), +]) +async def test_sendmsg(expected, message, send_headers, _content_length, start_byte): + writer = AsyncTestStreamWriter() + await sendmsg(message, writer, send_headers = send_headers, _content_length = _content_length, start_byte = start_byte) + assert writer.content == expected \ No newline at end of file diff --git a/tests/unit/jeb_utils/test_unpack_fields.py b/tests/unit/jeb_utils/test_unpack_fields.py new file mode 100644 index 0000000..e5ccb6b --- /dev/null +++ b/tests/unit/jeb_utils/test_unpack_fields.py @@ -0,0 +1,21 @@ +from src.jeb_utils.jebp_utils import unpack_fields +import pytest + +@pytest.mark.parametrize('packed,expected', [ + (b'', []), + (b'\x02\x00\x00', [b'']), + (b'\x01\x00', [b'']), + (b'\x01\x00\x01\x00', [b'', b'']), + (b'\x01\x00\x01\x01a', [b'', b'a']), + (b'\x01\x00\x01\x02ab', [b'', b'ab']), + (b'\x01\x01a\x01\x02ab', [b'a', b'ab']), + (b'\x01\x01a\x01\x02ab\x01\x00\x01\x0dHello, World!', [b'a', b'ab', b'', b'Hello, World!']), + (b'\x01\xff' + b'a' * 255, [b'a' * 255]), + ((b'\x01\xff' + b'a' * 255) * 2, [b'a' * 255, b'a' * 255]), + (b'\x02\x01\x00' + b'a' * 256, [b'a' * 256]), + ((b'\x02\x01\x00' + b'a' * 256) * 2, [b'a' * 256, b'a' * 256]), + (b'\x02\x01\x01' + b'a' * 257, [b'a' * 257]), + (b'\x02\x01\x01' + b'a' * 257 + b'\x01\x03abc', [b'a' * 257, b'abc']), +]) +def test_unpack_fields(packed, expected): + assert unpack_fields(packed) == expected \ No newline at end of file