From a45e5126978d72da231155af70ea1b38ee800e7e Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Motiejus=20Jak=C5=A1tys?= <motiejus@jakstys.lt>
Date: Mon, 24 May 2021 00:11:58 +0300
Subject: [PATCH] byteCounter tests

---
 internal/bytecounter/bytecounter_test.go | 31 ++++++++++++++++++------
 rootfs/BUILD                             |  5 +++-
 rootfs/rootfs.go                         |  5 ++--
 3 files changed, 31 insertions(+), 10 deletions(-)

diff --git a/internal/bytecounter/bytecounter_test.go b/internal/bytecounter/bytecounter_test.go
index 17f5c81..000c548 100644
--- a/internal/bytecounter/bytecounter_test.go
+++ b/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])
 }
diff --git a/rootfs/BUILD b/rootfs/BUILD
index 978f0b9..9c34769 100644
--- a/rootfs/BUILD
+++ b/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(
diff --git a/rootfs/rootfs.go b/rootfs/rootfs.go
index 8fd2e5d..14f0550 100644
--- a/rootfs/rootfs.go
+++ b/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