Nick Craig-Wood
7744f98960
NoCloserNotify hides the Close method of the reader passed in like NoCloser, but calls a notify function (once only) when the returned body is closed. This lets callers of http.NewRequest find out when the transport has finished with a request body, as it is documented to possibly close it in a different goroutine after the request has finished.
105 lines
2.1 KiB
Go
105 lines
2.1 KiB
Go
package readers
|
|
|
|
import (
|
|
"errors"
|
|
"io"
|
|
"testing"
|
|
|
|
"github.com/stretchr/testify/assert"
|
|
)
|
|
|
|
var errRead = errors.New("read error")
|
|
|
|
type readOnly struct{}
|
|
|
|
func (readOnly) Read(p []byte) (n int, err error) {
|
|
return 0, io.EOF
|
|
}
|
|
|
|
type readClose struct{}
|
|
|
|
func (readClose) Read(p []byte) (n int, err error) {
|
|
return 0, errRead
|
|
}
|
|
|
|
func (readClose) Close() (err error) {
|
|
return io.EOF
|
|
}
|
|
|
|
type readCloseWriteTo struct {
|
|
readClose
|
|
}
|
|
|
|
func (readCloseWriteTo) WriteTo(w io.Writer) (n int64, err error) {
|
|
return 42, errRead
|
|
}
|
|
|
|
func TestNoCloser(t *testing.T) {
|
|
assert.Equal(t, nil, NoCloser(nil))
|
|
|
|
ro := readOnly{}
|
|
assert.Equal(t, ro, NoCloser(ro))
|
|
|
|
rc := readClose{}
|
|
nc := NoCloser(rc)
|
|
assert.NotEqual(t, nc, rc)
|
|
|
|
_, hasClose := nc.(io.Closer)
|
|
assert.False(t, hasClose)
|
|
|
|
_, err := nc.Read(nil)
|
|
assert.Equal(t, errRead, err)
|
|
|
|
_, hasWriteTo := nc.(io.WriterTo)
|
|
assert.False(t, hasWriteTo)
|
|
|
|
rcw := readCloseWriteTo{}
|
|
ncw := NoCloser(rcw)
|
|
assert.NotEqual(t, ncw, rcw)
|
|
|
|
_, hasClose = ncw.(io.Closer)
|
|
assert.False(t, hasClose)
|
|
|
|
wt, hasWriteTo := ncw.(io.WriterTo)
|
|
assert.True(t, hasWriteTo)
|
|
n, err := wt.WriteTo(nil)
|
|
assert.Equal(t, int64(42), n)
|
|
assert.Equal(t, errRead, err)
|
|
}
|
|
|
|
func TestNoCloserNotify(t *testing.T) {
|
|
assert.Nil(t, NoCloserNotify(nil, func() {}))
|
|
|
|
notified := 0
|
|
notify := func() { notified++ }
|
|
|
|
// A reader without Close or WriteTo
|
|
nc := NoCloserNotify(readOnly{}, notify)
|
|
_, err := nc.Read(nil)
|
|
assert.Equal(t, io.EOF, err)
|
|
_, hasWriteTo := nc.(io.WriterTo)
|
|
assert.False(t, hasWriteTo)
|
|
|
|
// Close doesn't close the underlying reader (whose Close returns
|
|
// io.EOF) and notifies once only
|
|
notified = 0
|
|
nc = NoCloserNotify(readClose{}, notify)
|
|
_, err = nc.Read(nil)
|
|
assert.Equal(t, errRead, err)
|
|
assert.NoError(t, nc.Close())
|
|
assert.Equal(t, 1, notified)
|
|
assert.NoError(t, nc.Close())
|
|
assert.Equal(t, 1, notified)
|
|
|
|
// WriteTo is forwarded
|
|
notified = 0
|
|
ncw := NoCloserNotify(readCloseWriteTo{}, notify)
|
|
wt, hasWriteTo := ncw.(io.WriterTo)
|
|
assert.True(t, hasWriteTo)
|
|
n, err := wt.WriteTo(nil)
|
|
assert.Equal(t, int64(42), n)
|
|
assert.Equal(t, errRead, err)
|
|
assert.NoError(t, ncw.Close())
|
|
assert.Equal(t, 1, notified)
|
|
}
|