stream_test.go

package netpipe

import (
	"bytes"
	"net"
	"testing"
)

// TestStreamRoundTrip proves a large payload chunked and reassembled correctly.
func TestStreamRoundTrip(t *testing.T) {
	server, client := net.Pipe()
	defer server.Close()
	defer client.Close()

	// 256 KB payload, 64 KB chunks → 4 chunks
	msg := make([]byte, 256*1024)
	for i := range msg {
		msg[i] = byte(i % 256)
	}

	go func() {
		if err := writeStream(client, msg, DefaultChunkSize); err != nil {
			t.Errorf("writeStream: %v", err)
		}
	}()

	buf := newStreamBuffer()
	for {
		flags, body, err := readFrameFull(server)
		if err != nil {
			t.Fatalf("readFrameFull: %v", err)
		}

		if flags&flagStream == 0 {
			t.Fatal("expected flagStream to be set")
		}

		streamID, index, total, chunk, err := parseStreamBody(body)
		if err != nil {
			t.Fatalf("parseStreamBody: %v", err)
		}

		assembled, complete := buf.addChunk(streamID, index, total, chunk)
		if complete {
			if !bytes.Equal(assembled, msg) {
				t.Fatalf("reassembled data mismatch (got %d bytes, want %d)", len(assembled), len(msg))
			}
			return // success
		}
	}
}

// TestStreamSmallPayload proves streaming works for a payload smaller than one chunk.
func TestStreamSmallPayload(t *testing.T) {
	server, client := net.Pipe()
	defer server.Close()
	defer client.Close()

	msg := []byte("tiny payload")

	go func() {
		writeStream(client, msg, DefaultChunkSize)
	}()

	flags, body, err := readFrameFull(server)
	if err != nil {
		t.Fatalf("readFrameFull: %v", err)
	}
	if flags&flagStream == 0 {
		t.Fatal("expected flagStream")
	}

	streamID, index, total, chunk, err := parseStreamBody(body)
	if err != nil {
		t.Fatalf("parseStreamBody: %v", err)
	}

	if total != 1 {
		t.Fatalf("expected 1 total chunk, got %d", total)
	}

	assembled, complete := newStreamBuffer().addChunk(streamID, index, total, chunk)
	if !complete {
		t.Fatal("expected complete after single chunk")
	}
	if !bytes.Equal(assembled, msg) {
		t.Fatalf("got %q, want %q", assembled, msg)
	}
}

// TestStreamDuplicateChunkIgnored proves duplicate chunks don't corrupt reassembly.
func TestStreamDuplicateChunkIgnored(t *testing.T) {
	buf := newStreamBuffer()

	chunk0 := []byte("aaaa")
	chunk1 := []byte("bbbb")

	// add chunk 0
	_, complete := buf.addChunk("stream-1", 0, 2, chunk0)
	if complete {
		t.Fatal("should not be complete yet")
	}

	// duplicate chunk 0 - should be ignored
	_, complete = buf.addChunk("stream-1", 0, 2, chunk0)
	if complete {
		t.Fatal("duplicate should not trigger completion")
	}

	// add chunk 1 - should complete
	assembled, complete := buf.addChunk("stream-1", 1, 2, chunk1)
	if !complete {
		t.Fatal("should be complete now")
	}

	expected := append(chunk0, chunk1...)
	if !bytes.Equal(assembled, expected) {
		t.Fatalf("got %q, want %q", assembled, expected)
	}
}

// TestStreamBufferCleanup proves completed streams are removed from the buffer.
func TestStreamBufferCleanup(t *testing.T) {
	buf := newStreamBuffer()

	buf.addChunk("s1", 0, 1, []byte("data"))

	buf.mu.Lock()
	_, exists := buf.streams["s1"]
	buf.mu.Unlock()

	if exists {
		t.Fatal("completed stream should be cleaned up from buffer")
	}
}