45 lines
1.0 KiB
Go
45 lines
1.0 KiB
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))
|
|
}
|