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])
 }

rootfs/BUILD

@@ -9,7 +9,10 @@ go_library(
     ],
     importpath = "github.com/motiejus/code/undocker/rootfs",
     visibility = ["//visibility:public"],
-    deps = ["@org_uber_go_multierr//:go_default_library"],
+    deps = [
+        "//src/undocker/internal/bytecounter:go_default_library",
+        "@org_uber_go_multierr//:go_default_library",
+    ],
 )
 
 go_test(

rootfs/rootfs.go

@@ -9,6 +9,7 @@ import (
 	"path/filepath"
 	"strings"
 
+     "github.com/motiejus/code/undocker/internal/bytecounter"
 	"go.uber.org/multierr"
 )
 
@@ -39,12 +40,12 @@ func New(rd io.ReadSeeker) *RootFS {
 
 // WriteTo writes a docker image to an open tarball.
 func (r *RootFS) WriteTo(w io.Writer) (n int64, err error) {
-	wr := &byteCounter{rw: w}
+	wr := bytecounter.New(w)
 	tr := tar.NewReader(r.rd)
 	tw := tar.NewWriter(wr)
 	defer func() {
 		err = multierr.Append(err, tw.Close())
-		n = wr.n
+		n = wr.N
 	}()
 
 	// layerOffsets maps a layer name (a9b123c0daa/layer.tar) to it's offset