add io helper

Author: eli-darkly

io.go

@@ -0,0 +1,55 @@
+package helpers
+
+import (
+	"io"
+	"time"
+)
+
+// ReadWithTimeout reads data until it gets the desired number of bytes or times out.
+//
+// This is an inefficient implementation that should only be used in tests.
+func ReadWithTimeout(r io.Reader, n int, timeout time.Duration) []byte {
+	byteCh := make(chan byte)
+	closer := make(chan struct{})
+
+	go func() {
+		count := 0
+		for {
+			select {
+			case <-closer:
+				return
+			default:
+				b := make([]byte, 1)
+				got, err := r.Read(b)
+				if err != nil {
+					return
+				}
+				if got > 0 {
+					byteCh <- b[0]
+				}
+				count++
+				if count >= n {
+					return
+				}
+			}
+		}
+	}()
+
+	buf := make([]byte, 0, n)
+	deadline := time.After(timeout)
+
+ReadLoop:
+	for {
+		select {
+		case b := <-byteCh:
+			buf = append(buf, b)
+			if len(buf) >= n {
+				break ReadLoop
+			}
+		case <-deadline:
+			break ReadLoop
+		}
+	}
+	close(closer)
+	return buf
+}

io_test.go

@@ -0,0 +1,43 @@
+package helpers
+
+import (
+	"bytes"
+	"io"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestReadWithTimeout(t *testing.T) {
+	var buf1 bytes.Buffer
+	data := ReadWithTimeout(&buf1, 5, time.Millisecond*50)
+	assert.Len(t, data, 0)
+
+	var buf2 bytes.Buffer
+	buf2.WriteString("he")
+	data = ReadWithTimeout(&buf2, 5, time.Millisecond*50)
+	assert.Equal(t, "he", string(data))
+
+	var buf3 bytes.Buffer
+	buf3.WriteString("hello")
+	data = ReadWithTimeout(&buf3, 5, time.Millisecond*50)
+	assert.Equal(t, "hello", string(data))
+
+	r1, w1 := io.Pipe()
+	go func() {
+		w1.Write([]byte("good"))
+		time.Sleep(10 * time.Millisecond)
+		w1.Write([]byte("bye"))
+	}()
+	data = ReadWithTimeout(r1, 7, time.Millisecond*100)
+	assert.Equal(t, "goodbye", string(data))
+
+	r2, w2 := io.Pipe()
+	go func() {
+		w2.Write([]byte("good"))
+	}()
+	data = ReadWithTimeout(r2, 7, time.Millisecond*100)
+	time.Sleep(time.Millisecond * 100)
+	assert.Equal(t, "good", string(data))
+}