byteCounter tests

internal/bytecounter/bytecounter_test.go

@@ -4,7 +4,6 @@ import (
 	"bytes"
 	"io"
 	"testing"
-	"testing/iotest"
 
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
@@ -13,12 +12,30 @@ import (
 func TestByteCounter(t *testing.T) {
 	r := bytes.NewBufferString("0123456789")
 	w := bytes.NewBuffer(nil)
-
-	tw := iotest.TruncateWriter(w, 4)
-	bc := New(tw)
-
+	bc := New(w)
 	_, err := io.Copy(bc, r)
 	require.NoError(t, err)
-	assert.Len(t, w.Bytes(), 4)
-	assert.Equal(t, 4, bc.N)
+	assert.Equal(t, []byte("0123456789"), w.Bytes())
+	assert.Equal(t, int64(10), bc.N)
+}
+
+func TestSingleByteCounter(t *testing.T) {
+	r := bytes.NewBufferString("0123456789")
+	w := bytes.NewBuffer(nil)
+	tw := &shortWriter{w}
+	bc := New(tw)
+	_, err := io.Copy(bc, r)
+	assert.EqualError(t, err, "short write")
+	assert.Len(t, w.Bytes(), 1)
+	assert.Equal(t, int64(1), bc.N)
+}
+
+// shortWriter writes only first byte
+type shortWriter struct {
+	w io.Writer
+}
+
+// Write writes a byte to the underlying writer
+func (f *shortWriter) Write(p []byte) (int, error) {
+	return f.w.Write(p[0:1])
 }