namespace MoShiro.Protocol.Tests;

public sealed class TcpWireTests
{
    [Theory]
    [InlineData(0)]
    [InlineData(64)]
    [InlineData(16384)]
    public async Task Message_UsesOneWriteForPrefixAndBody_AndPreservesFraming(int size)
    {
        using var stream = new CountingStream();
        byte[] body = Enumerable.Range(0, size).Select(i => (byte)i).ToArray();
        await TcpWire.WriteMessageAsync(stream, body, default);
        Assert.Equal(1, stream.Writes);
        stream.Position = 0;
        Assert.Equal(body, await TcpWire.ReadMessageAsync(stream, default));
        Assert.Null(await TcpWire.ReadMessageAsync(stream, default));
    }

    private sealed class CountingStream : MemoryStream
    {
        public int Writes { get; private set; }

        public override ValueTask WriteAsync(ReadOnlyMemory<byte> buffer, CancellationToken cancellationToken = default)
        {
            Writes++;
            return base.WriteAsync(buffer, cancellationToken);
        }
    }
}