nnet/internal/unpacker/length_field_test.go

package unpacker

import (
	"testing"

	unpackerpkg "git.noahlan.cn/noahlan/nnet/v2/pkg/unpacker"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
)

func TestLengthFieldUnpacker(t *testing.T) {
	config := unpackerpkg.LengthFieldUnpacker{
		LengthFieldOffset:   0,
		LengthFieldLength:   4,
		LengthAdjustment:    0,
		InitialBytesToStrip: 4, // 跳过长度字段，只返回数据部分
		MaxBufferSize:       1024,
	}
	u := NewLengthFieldUnpacker(config)

	// 创建测试数据：4字节长度 + 数据
	data := make([]byte, 8)
	data[0] = 0
	data[1] = 0
	data[2] = 0
	data[3] = 4 // 长度 = 4
	copy(data[4:], []byte("test"))

	messages, remaining, consumed, err := u.Unpack(data)
	require.NoError(t, err, "Expected no error")
	assert.Equal(t, 1, len(messages), "Expected 1 message")
	assert.Equal(t, "test", string(messages[0]), "Expected 'test'")
	assert.Equal(t, 0, len(remaining), "Expected no remaining data")
	assert.Equal(t, len(data), consumed, "Expected consumed equals input data length")
}

func TestLengthFieldUnpackerMultipleMessages(t *testing.T) {
	config := unpackerpkg.LengthFieldUnpacker{
		LengthFieldOffset:   0,
		LengthFieldLength:   4,
		LengthAdjustment:    0,
		InitialBytesToStrip: 4, // 跳过长度字段，只返回数据部分
		MaxBufferSize:       1024,
	}
	u := NewLengthFieldUnpacker(config)

	// 创建多个消息
	data := make([]byte, 16)
	// 第一个消息：长度4，数据"test"
	data[0] = 0
	data[1] = 0
	data[2] = 0
	data[3] = 4
	copy(data[4:], []byte("test"))
	// 第二个消息：长度4，数据"data"
	data[8] = 0
	data[9] = 0
	data[10] = 0
	data[11] = 4
	copy(data[12:], []byte("data"))

	messages, _, consumed, err := u.Unpack(data)
	require.NoError(t, err, "Expected no error")
	assert.Equal(t, 2, len(messages), "Expected 2 messages")
	assert.Equal(t, "test", string(messages[0]), "Expected 'test'")
	assert.Equal(t, "data", string(messages[1]), "Expected 'data'")
	assert.Equal(t, len(data), consumed, "Expected consumed equals input data length")
}

func TestLengthFieldUnpackerIncompleteMessage(t *testing.T) {
	config := unpackerpkg.LengthFieldUnpacker{
		LengthFieldOffset:   0,
		LengthFieldLength:   4,
		LengthAdjustment:    0,
		InitialBytesToStrip: 0,
		MaxBufferSize:       1024,
	}
	u := NewLengthFieldUnpacker(config)

	// 测试不完整的消息（只有长度字段）
	data := []byte{0, 0, 0, 4}
	messages, remaining, consumed, err := u.Unpack(data)
	require.NoError(t, err, "Expected no error")
	assert.Equal(t, 0, len(messages), "Expected 0 messages")
	assert.Equal(t, 4, len(remaining), "Expected 4 bytes remaining")
	assert.Equal(t, len(data), consumed, "Expected all data consumed when no complete message")
}

func TestLengthFieldUnpackerConsumesBufferedRemainder(t *testing.T) {
	config := unpackerpkg.LengthFieldUnpacker{
		LengthFieldOffset:   0,
		LengthFieldLength:   1,
		LengthAdjustment:    0,
		InitialBytesToStrip: 0,
		MaxBufferSize:       1024,
	}
	u := NewLengthFieldUnpacker(config)

	messages, remaining, consumed, err := u.Unpack([]byte{3, 'a'})
	require.NoError(t, err)
	assert.Empty(t, messages)
	assert.Equal(t, []byte{3, 'a'}, remaining)
	assert.Equal(t, 2, consumed)

	messages, remaining, consumed, err = u.Unpack([]byte{'b', 'c', 1})
	require.NoError(t, err)
	require.Len(t, messages, 1)
	assert.Equal(t, []byte{3, 'a', 'b', 'c'}, messages[0])
	assert.Equal(t, []byte{1}, remaining)
	assert.Equal(t, 3, consumed, "Expected complete and buffered input bytes to be consumed")
}

func TestLengthFieldUnpackerRejectsInvalidStrip(t *testing.T) {
	config := unpackerpkg.LengthFieldUnpacker{
		LengthFieldOffset:   0,
		LengthFieldLength:   1,
		LengthAdjustment:    0,
		InitialBytesToStrip: 5,
		MaxBufferSize:       1024,
	}
	u := NewLengthFieldUnpacker(config)

	_, _, _, err := u.Unpack([]byte{1, 'a'})
	assert.Error(t, err)
}

func TestLengthFieldUnpackerPack(t *testing.T) {
	config := unpackerpkg.LengthFieldUnpacker{
		LengthFieldOffset:   0,
		LengthFieldLength:   4,
		LengthAdjustment:    0,
		InitialBytesToStrip: 0,
		MaxBufferSize:       1024,
	}
	u := NewLengthFieldUnpacker(config)

	// 测试打包
	data := []byte("test")
	packed, err := u.Pack(data)
	require.NoError(t, err, "Expected no error")
	assert.Equal(t, 8, len(packed), "Expected 8 bytes")
	assert.Equal(t, byte(4), packed[3], "Expected length 4")
	assert.Equal(t, "test", string(packed[4:]), "Expected 'test'")
}

func TestLengthFieldUnpackerBufferSizeLimit(t *testing.T) {
	// 测试缓冲区大小限制
	config := unpackerpkg.LengthFieldUnpacker{
		LengthFieldOffset:   0,
		LengthFieldLength:   4,
		LengthAdjustment:    0,
		InitialBytesToStrip: 0,
		MaxBufferSize:       10, // 小缓冲区
	}
	u := NewLengthFieldUnpacker(config)

	// 创建超过缓冲区大小的数据
	data := make([]byte, 20)
	data[0] = 0
	data[1] = 0
	data[2] = 0
	data[3] = 20 // 长度 = 20

	_, _, _, err := u.Unpack(data)
	assert.Error(t, err, "Expected error for buffer size exceeded")
}