bytecounter

internal/bytecounter/BUILD

@@ -0,0 +1,18 @@
+load("@io_bazel_rules_go//go:def.bzl", "go_library", "go_test")
+
+go_library(
+    name = "go_default_library",
+    srcs = ["bytecounter.go"],
+    importpath = "github.com/motiejus/code/undocker/internal/bytecounter",
+    visibility = ["//src/undocker:__subpackages__"],
+)
+
+go_test(
+    name = "go_default_test",
+    srcs = ["bytecounter_test.go"],
+    embed = [":go_default_library"],
+    deps = [
+        "@com_github_stretchr_testify//assert:go_default_library",
+        "@com_github_stretchr_testify//require:go_default_library",
+    ],
+)

internal/bytecounter/bytecounter.go

@@ -0,0 +1,20 @@
+package bytecounter
+
+import "io"
+
+// ByteCounter is an io.Writer that counts bytes written to it
+type ByteCounter struct {
+	N int64
+	w io.Writer
+}
+
+// New returns a new ByteCounter
+func New(w io.Writer) *ByteCounter {
+	return &ByteCounter{w: w}
+}
+
+// Write writes to the underlying io.Writer and counts total written bytes
+func (b *ByteCounter) Write(data []byte) (n int, err error) {
+	defer func() { b.N += int64(n) }()
+	return b.w.Write(data)
+}

internal/bytecounter/bytecounter_test.go

@@ -0,0 +1,24 @@
+package bytecounter
+
+import (
+	"bytes"
+	"io"
+	"testing"
+	"testing/iotest"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+func TestByteCounter(t *testing.T) {
+	r := bytes.NewBufferString("0123456789")
+	w := bytes.NewBuffer(nil)
+
+	tw := iotest.TruncateWriter(w, 4)
+	bc := New(tw)
+
+	_, err := io.Copy(bc, r)
+	require.NoError(t, err)
+	assert.Len(t, w.Bytes(), 4)
+	assert.Equal(t, 4, bc.N)
+}