move implementation to internal/ directory (#828)
This commit is contained in:
committed by
GitHub
parent
b9b9ad10cf
commit
908807bd59
@@ -0,0 +1,36 @@
|
||||
package limitio
|
||||
|
||||
import "io"
|
||||
|
||||
type readCloser struct {
|
||||
read int64
|
||||
limit int64
|
||||
r io.ReadCloser
|
||||
}
|
||||
|
||||
func ReadCloser(rc io.ReadCloser, limit int64) io.ReadCloser {
|
||||
return &readCloser{0, limit, rc}
|
||||
}
|
||||
|
||||
var _ io.ReadCloser = (*readCloser)(nil)
|
||||
|
||||
func (r *readCloser) Read(b []byte) (int, error) {
|
||||
|
||||
if len(b) == 0 {
|
||||
return 0, nil
|
||||
}
|
||||
|
||||
if r.read == r.limit {
|
||||
return 0, io.EOF
|
||||
}
|
||||
|
||||
if r.read+int64(len(b)) >= r.limit {
|
||||
b = b[:int(r.limit-r.read)]
|
||||
}
|
||||
|
||||
readN, err := r.r.Read(b)
|
||||
r.read += int64(readN)
|
||||
return readN, err
|
||||
}
|
||||
|
||||
func (r *readCloser) Close() error { return r.r.Close() }
|
||||
@@ -0,0 +1,44 @@
|
||||
package limitio
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"io"
|
||||
"testing"
|
||||
|
||||
"github.com/stretchr/testify/assert"
|
||||
)
|
||||
|
||||
type mockRC struct {
|
||||
r io.Reader
|
||||
closed bool
|
||||
}
|
||||
|
||||
func newMockRC(r io.Reader) *mockRC { return &mockRC{r, false} }
|
||||
|
||||
func (m mockRC) Read(b []byte) (int, error) {
|
||||
if m.closed {
|
||||
return 0, io.ErrClosedPipe
|
||||
}
|
||||
return m.r.Read(b)
|
||||
}
|
||||
|
||||
func (m *mockRC) Close() error {
|
||||
m.closed = true
|
||||
return nil
|
||||
}
|
||||
|
||||
func TestReadCloser(t *testing.T) {
|
||||
foobarReader := bytes.NewReader([]byte("foobar2342"))
|
||||
mock := newMockRC(foobarReader)
|
||||
limited := ReadCloser(mock, 6)
|
||||
var buf [20]byte
|
||||
n, err := limited.Read(buf[:])
|
||||
assert.NoError(t, err)
|
||||
assert.Equal(t, 6, n)
|
||||
assert.Equal(t, buf[:n], []byte("foobar"))
|
||||
n, err = limited.Read(buf[:])
|
||||
assert.Equal(t, 0, n)
|
||||
assert.Equal(t, io.EOF, err)
|
||||
limited.Close()
|
||||
assert.True(t, mock.closed)
|
||||
}
|
||||
Reference in New Issue
Block a user