Use iotest.OneByteReader instead of self implemented one

Client: go

This is a trivial unit test improvement from the last commit. Of course
I only discovered the existence of testing/iotest package after I
re-invented the wheel.

2 files changed, 8 insertions, 19 deletions

diff --git a/lib/go/thrift/framed_transport_test.go b/lib/go/thrift/framed_transport_test.go
index 4e7d9cae0..d23ec595e 100644
--- a/lib/go/thrift/framed_transport_test.go
+++ b/lib/go/thrift/framed_transport_test.go
@@ -24,6 +24,7 @@ import (
 	"io"
 	"strings"
 	"testing"
+	"testing/iotest"
 )
 
 func TestFramedTransport(t *testing.T) {
@@ -51,7 +52,7 @@ func TestTFramedTransportReuseTransport(t *testing.T) {
 			}
 
 			// read
-			read, err := io.ReadAll(oneAtATimeReader{reader})
+			read, err := io.ReadAll(iotest.OneByteReader(reader))
 			if err != nil {
 				t.Errorf("Failed to read on #%d: %v", i, err)
 			}
@@ -80,9 +81,9 @@ func TestTFramedTransportReuseTransport(t *testing.T) {
 			var buf []byte
 			var err error
 			if i%2 == 0 {
-				// on even calls, use oneAtATimeReader to make
+				// on even calls, use OneByteReader to make
 				// sure that small reads are fine
-				buf, err = io.ReadAll(io.LimitReader(oneAtATimeReader{reader}, int64(size)))
+				buf, err = io.ReadAll(io.LimitReader(iotest.OneByteReader(reader), int64(size)))
 			} else {
 				// on odd calls, make sure that we don't read
 				// more than written per frame
diff --git a/lib/go/thrift/header_transport_test.go b/lib/go/thrift/header_transport_test.go
index 44d0284db..125a5fd8c 100644
--- a/lib/go/thrift/header_transport_test.go
+++ b/lib/go/thrift/header_transport_test.go
@@ -25,6 +25,7 @@ import (
 	"io"
 	"strings"
 	"testing"
+	"testing/iotest"
 	"testing/quick"
 )
 
@@ -325,7 +326,7 @@ func TestTHeaderTransportReuseTransport(t *testing.T) {
 			}
 
 			// read
-			read, err := io.ReadAll(oneAtATimeReader{reader})
+			read, err := io.ReadAll(iotest.OneByteReader(reader))
 			if err != nil {
 				t.Errorf("Failed to read on #%d: %v", i, err)
 			}
@@ -354,9 +355,9 @@ func TestTHeaderTransportReuseTransport(t *testing.T) {
 			var buf []byte
 			var err error
 			if i%2 == 0 {
-				// on even calls, use oneAtATimeReader to make
+				// on even calls, use OneByteReader to make
 				// sure that small reads are fine
-				buf, err = io.ReadAll(io.LimitReader(oneAtATimeReader{reader}, int64(size)))
+				buf, err = io.ReadAll(io.LimitReader(iotest.OneByteReader(reader), int64(size)))
 			} else {
 				// on odd calls, make sure that we don't read
 				// more than written per frame
@@ -374,16 +375,3 @@ func TestTHeaderTransportReuseTransport(t *testing.T) {
 		}
 	})
 }
-
-type oneAtATimeReader struct {
-	io.Reader
-}
-
-// oneAtATimeReader forces every Read call to only read 1 byte out,
-// thus forces the underlying reader's Read to be called multiple times.
-func (o oneAtATimeReader) Read(buf []byte) (int, error) {
-	if len(buf) < 1 {
-		return o.Reader.Read(buf)
-	}
-	return o.Reader.Read(buf[:1])
-}