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