infra/proxyd/reader_test.go

package proxyd

import (
	"io"
	"strings"
	"testing"

	"github.com/stretchr/testify/require"
)

func TestLimitReader(t *testing.T) {
	data := "hellohellohellohello"
	r := LimitReader(strings.NewReader(data), 3)
	buf := make([]byte, 3)

	// Buffer reads OK
	n, err := r.Read(buf)
	require.NoError(t, err)
	require.Equal(t, 3, n)

	// Buffer is over limit
	n, err = r.Read(buf)
	require.Equal(t, ErrLimitReaderOverLimit, err)
	require.Equal(t, 0, n)

	// Buffer on initial read is over size
	buf = make([]byte, 16)
	r = LimitReader(strings.NewReader(data), 3)
	n, err = r.Read(buf)
	require.NoError(t, err)
	require.Equal(t, 3, n)

	// test with read all where the limit is less than the data
	r = LimitReader(strings.NewReader(data), 3)
	out, err := io.ReadAll(r)
	require.Equal(t, ErrLimitReaderOverLimit, err)
	require.Equal(t, "hel", string(out))

	// test with read all where the limit is more than the data
	r = LimitReader(strings.NewReader(data), 21)
	out, err = io.ReadAll(r)
	require.NoError(t, err)
	require.Equal(t, data, string(out))
}