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)