acdream/tests/AcDream.Core.Net.Tests/Packets/FragmentAssemblerTests.cs

using AcDream.Core.Net.Packets;

namespace AcDream.Core.Net.Tests.Packets;

public class FragmentAssemblerTests
{
    private static MessageFragment MakeFrag(uint id, ushort count, ushort index, byte[] payload, ushort queue = 7)
        => new(
            new MessageFragmentHeader
            {
                Id        = id,
                Count     = count,
                Index     = index,
                TotalSize = (ushort)(MessageFragmentHeader.Size + payload.Length),
                Queue     = queue,
            },
            payload);

    [Fact]
    public void Ingest_SingleFragmentMessage_ReleasesImmediately()
    {
        var assembler = new FragmentAssembler();
        var frag = MakeFrag(id: 1, count: 1, index: 0, payload: new byte[] { 1, 2, 3 }, queue: 42);

        var result = assembler.Ingest(frag, out var queue);

        Assert.NotNull(result);
        Assert.Equal(new byte[] { 1, 2, 3 }, result);
        Assert.Equal(42, queue);
        Assert.Equal(0, assembler.PartialCount);
    }

    [Fact]
    public void Ingest_ThreeFragmentsInOrder_ReleasesOnLast()
    {
        // Queue is a property of the logical message, not individual fragments,
        // so all three fragments carry the same queue value (captured from the
        // first arrival). Testing with queue=9 on all three.
        var assembler = new FragmentAssembler();

        Assert.Null(assembler.Ingest(MakeFrag(7, 3, 0, new byte[] { 0xAA, 0xBB }, queue: 9), out _));
        Assert.Equal(1, assembler.PartialCount);
        Assert.Null(assembler.Ingest(MakeFrag(7, 3, 1, new byte[] { 0xCC, 0xDD }, queue: 9), out _));
        var result = assembler.Ingest(MakeFrag(7, 3, 2, new byte[] { 0xEE }, queue: 9), out var queue);

        Assert.NotNull(result);
        Assert.Equal(new byte[] { 0xAA, 0xBB, 0xCC, 0xDD, 0xEE }, result);
        Assert.Equal(9, queue);
        Assert.Equal(0, assembler.PartialCount);
    }

    [Fact]
    public void Ingest_OutOfOrderFragments_ReleasesCorrectlyOnLastArrival()
    {
        // Arrive as index 2, then 0, then 1 — the last arrival (index 1) is
        // neither the first nor the last index, so this tests that the
        // assembler releases on "count full", not "last index".
        var assembler = new FragmentAssembler();

        Assert.Null(assembler.Ingest(MakeFrag(3, 3, 2, new byte[] { 0xCC }), out _));
        Assert.Null(assembler.Ingest(MakeFrag(3, 3, 0, new byte[] { 0xAA }), out _));
        var result = assembler.Ingest(MakeFrag(3, 3, 1, new byte[] { 0xBB }), out _);

        Assert.NotNull(result);
        // Result must be assembled in INDEX order, not arrival order.
        Assert.Equal(new byte[] { 0xAA, 0xBB, 0xCC }, result);
    }

    [Fact]
    public void Ingest_DuplicateFragment_IsIdempotent()
    {
        var assembler = new FragmentAssembler();
        Assert.Null(assembler.Ingest(MakeFrag(5, 2, 0, new byte[] { 0x11 }), out _));
        // Resend index 0 — should not double-count or corrupt state.
        Assert.Null(assembler.Ingest(MakeFrag(5, 2, 0, new byte[] { 0x11 }), out _));
        // Assembler should still be waiting for index 1.
        Assert.Equal(1, assembler.PartialCount);

        var result = assembler.Ingest(MakeFrag(5, 2, 1, new byte[] { 0x22 }), out _);
        Assert.NotNull(result);
        Assert.Equal(new byte[] { 0x11, 0x22 }, result);
    }

    [Fact]
    public void Ingest_MissingFragment_DoesNotRelease()
    {
        var assembler = new FragmentAssembler();
        Assert.Null(assembler.Ingest(MakeFrag(9, 3, 0, new byte[] { 1 }), out _));
        Assert.Null(assembler.Ingest(MakeFrag(9, 3, 2, new byte[] { 3 }), out _));
        // Only 2 of 3 arrived → still waiting
        Assert.Equal(1, assembler.PartialCount);
    }

    [Fact]
    public void Ingest_TwoIndependentMessages_BuiltInParallel()
    {
        var assembler = new FragmentAssembler();

        Assert.Null(assembler.Ingest(MakeFrag(100, 2, 0, new byte[] { 0xA1 }), out _));
        Assert.Null(assembler.Ingest(MakeFrag(200, 2, 0, new byte[] { 0xB1 }), out _));
        Assert.Equal(2, assembler.PartialCount);

        var resultA = assembler.Ingest(MakeFrag(100, 2, 1, new byte[] { 0xA2 }), out _);
        Assert.Equal(new byte[] { 0xA1, 0xA2 }, resultA);
        Assert.Equal(1, assembler.PartialCount);

        var resultB = assembler.Ingest(MakeFrag(200, 2, 1, new byte[] { 0xB2 }), out _);
        Assert.Equal(new byte[] { 0xB1, 0xB2 }, resultB);
        Assert.Equal(0, assembler.PartialCount);
    }

    [Fact]
    public void DropAll_ClearsInFlightPartials()
    {
        var assembler = new FragmentAssembler();
        assembler.Ingest(MakeFrag(1, 5, 0, new byte[] { 1 }), out _);
        assembler.Ingest(MakeFrag(2, 5, 0, new byte[] { 2 }), out _);
        Assert.Equal(2, assembler.PartialCount);

        assembler.DropAll();
        Assert.Equal(0, assembler.PartialCount);
    }
}