1 package client
2
3 import (
4 "bytes"
5 "io"
6 "testing"
7
8 "github.com/stretchr/testify/assert"
9 "github.com/stretchr/testify/require"
10 )
11
12 func newCountingReader(rdr io.Reader, readOnce bool) *countingReadCloser {
13 return &countingReadCloser{
14 rdr: rdr,
15 readOnce: readOnce,
16 }
17 }
18
19 type countingReadCloser struct {
20 rdr io.Reader
21 readOnce bool
22 readCalled int
23 closeCalled int
24 }
25
26 func (c *countingReadCloser) Read(b []byte) (int, error) {
27 c.readCalled++
28 if c.readCalled > 1 && c.readOnce {
29 return 0, io.EOF
30 }
31 return c.rdr.Read(b)
32 }
33
34 func (c *countingReadCloser) Close() error {
35 c.closeCalled++
36 return nil
37 }
38
39 func TestDrainingReadCloser(t *testing.T) {
40 rdr := newCountingReader(bytes.NewBufferString("There are many things to do"), false)
41 prevDisc := io.Discard
42 disc := bytes.NewBuffer(nil)
43 io.Discard = disc
44 defer func() { io.Discard = prevDisc }()
45
46 buf := make([]byte, 5)
47 ts := &drainingReadCloser{rdr: rdr}
48 _, err := ts.Read(buf)
49 require.NoError(t, err)
50 require.NoError(t, ts.Close())
51 assert.Equal(t, "There", string(buf))
52 assert.Equal(t, " are many things to do", disc.String())
53 assert.Equal(t, 3, rdr.readCalled)
54 assert.Equal(t, 1, rdr.closeCalled)
55 }
56
57 func TestDrainingReadCloser_SeenEOF(t *testing.T) {
58 rdr := newCountingReader(bytes.NewBufferString("There are many things to do"), true)
59 prevDisc := io.Discard
60 disc := bytes.NewBuffer(nil)
61 io.Discard = disc
62 defer func() { io.Discard = prevDisc }()
63
64 buf := make([]byte, 5)
65 ts := &drainingReadCloser{rdr: rdr}
66 _, err := ts.Read(buf)
67 require.NoError(t, err)
68 _, err = ts.Read(nil)
69 require.ErrorIs(t, err, io.EOF)
70 require.NoError(t, ts.Close())
71 assert.Equal(t, "There", string(buf))
72 assert.Empty(t, disc.String())
73 assert.Equal(t, 2, rdr.readCalled)
74 assert.Equal(t, 1, rdr.closeCalled)
75 }
76
View as plain text