50 lines
1.2 KiB
Go
50 lines
1.2 KiB
Go
package httpserver
|
|
|
|
import (
|
|
"bytes"
|
|
"io"
|
|
"testing"
|
|
|
|
"github.com/stretchr/testify/assert"
|
|
)
|
|
|
|
func TestLimitBufer(t *testing.T) {
|
|
buf := NewLimitBuffer(10)
|
|
n, err := buf.Write([]byte("1"))
|
|
assert.NoError(t, err)
|
|
assert.Equal(t, 1, n)
|
|
assert.Equal(t, "1", buf.String())
|
|
assert.Equal(t, 9, buf.Remain())
|
|
|
|
n, err = buf.Write([]byte("TestLimitBufer"))
|
|
// 设置limit但是不会返回 EOF
|
|
assert.NoError(t, err)
|
|
assert.Equal(t, 14, n)
|
|
assert.Equal(t, "1TestLimit...(total 15B)", buf.String())
|
|
assert.Equal(t, 0, buf.Remain())
|
|
|
|
n, err = buf.Write([]byte("TestLimitBufer"))
|
|
// 设置limit但是不会返回 EOF
|
|
assert.NoError(t, err)
|
|
assert.Equal(t, 14, n)
|
|
assert.Equal(t, "1TestLimit...(total 29B)", buf.String())
|
|
assert.Equal(t, 0, buf.Remain())
|
|
}
|
|
|
|
func TestTeeReader(t *testing.T) {
|
|
a := bytes.NewBuffer([]byte("TestLimitBufer"))
|
|
buf := NewLimitBuffer(10)
|
|
tee := TeeReadCloser(io.NopCloser(a), buf)
|
|
|
|
_, _ = io.CopyN(io.Discard, tee, 1)
|
|
assert.Equal(t, "T", buf.String())
|
|
assert.Equal(t, 9, buf.Remain())
|
|
|
|
io.Copy(io.Discard, io.LimitReader(tee, int64(buf.Remain())))
|
|
assert.Equal(t, "TestLimitB", buf.String())
|
|
assert.Equal(t, 0, buf.Remain())
|
|
|
|
io.Copy(io.Discard, tee)
|
|
assert.Equal(t, "TestLimitB...(total 14B)", buf.String())
|
|
}
|