From ef777d9e86e10e1c4ed946ff67c2bacf311bb790 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] 'want' can also accept different types

---
 rootfs/rootfs_test.go | 170 ++++++++++++++++++++++--------------------
 1 file changed, 89 insertions(+), 81 deletions(-)

diff --git a/rootfs/rootfs_test.go b/rootfs/rootfs_test.go
index da725f7..990356f 100644
--- a/rootfs/rootfs_test.go
+++ b/rootfs/rootfs_test.go
@@ -11,10 +11,91 @@ import (
 	"github.com/stretchr/testify/require"
 )
 
+func TestRootFS(t *testing.T) {
+	layer0 := tarball{
+		dir{name: "/", uid: 0},
+		file{name: "/file", uid: 0, contents: []byte("from 0")},
+	}
+
+	layer1 := tarball{
+		file{name: "/file", uid: 1, contents: []byte("from 1")},
+	}
+
+	layer2 := tarball{
+		dir{name: "/", uid: 2},
+	}
+
+	tests := []struct {
+		name    string
+		image   tarball
+		want    []extractable
+		wantErr string
+	}{
+		{
+			name:  "empty tarball",
+			image: tarball{manifest{}},
+			want:  []extractable{},
+		},
+		{
+			name:    "missing layer",
+			image:   tarball{manifest{"layer0/layer.tar"}},
+			wantErr: "bad or missing manifest.json",
+		},
+		{
+			name: "basic file overwrite, layer order mixed",
+			image: tarball{
+				file{name: "layer1/layer.tar", contents: layer1.bytes(t)},
+				file{name: "layer0/layer.tar", contents: layer0.bytes(t)},
+				manifest{"layer0/layer.tar", "layer1/layer.tar"},
+			},
+			want: []extractable{
+				dir{name: "/", uid: 0},
+				file{name: "/file", uid: 1, contents: []byte("from 1")},
+			},
+		},
+		{
+			name: "directory overwrite retains original dir",
+			image: tarball{
+				file{name: "layer2/layer.tar", contents: layer2.bytes(t)},
+				file{name: "layer0/layer.tar", contents: layer0.bytes(t)},
+				file{name: "layer1/layer.tar", contents: layer1.bytes(t)},
+				manifest{"layer0/layer.tar", "layer1/layer.tar", "layer2/layer.tar"},
+			},
+			want: []extractable{
+				dir{name: "/", uid: 0},
+				file{name: "/file", uid: 1, contents: []byte("from 1")},
+				dir{name: "/", uid: 2},
+			},
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			in := bytes.NewReader(tt.image.bytes(t))
+			out := bytes.Buffer{}
+
+			err := RootFS(in, &out)
+			if tt.wantErr != "" {
+				assert.EqualError(t, err, tt.wantErr)
+				return
+			}
+			require.NoError(t, err)
+			got := extract(t, &out)
+			assert.Equal(t, got, tt.want)
+		})
+	}
+}
+
+// Helpers
+
 type tarrable interface {
 	tar(*testing.T, *tar.Writer)
 }
 
+// extractable is an empty interface for comparing extracted outputs in tests.
+// Using that just to avoid the ugly `interface{}`.
+type extractable interface{}
+
 type dir struct {
 	name string
 	uid  int
@@ -64,7 +145,6 @@ type tarball []tarrable
 
 func (tb tarball) bytes(t *testing.T) []byte {
 	t.Helper()
-
 	buf := bytes.Buffer{}
 	tw := tar.NewWriter(&buf)
 	for _, member := range tb {
@@ -74,9 +154,9 @@ func (tb tarball) bytes(t *testing.T) []byte {
 	return buf.Bytes()
 }
 
-func extract(t *testing.T, f io.Reader) []file {
+func extract(t *testing.T, f io.Reader) []extractable {
 	t.Helper()
-	ret := []file{}
+	ret := []extractable{}
 	tr := tar.NewReader(f)
 	for {
 		hdr, err := tr.Next()
@@ -85,88 +165,16 @@ func extract(t *testing.T, f io.Reader) []file {
 		}
 		require.NoError(t, err)
 
-		elem := file{name: hdr.Name, uid: hdr.Uid}
-		if hdr.Typeflag == tar.TypeReg {
+		var elem extractable
+		switch hdr.Typeflag {
+		case tar.TypeDir:
+			elem = dir{name: hdr.Name, uid: hdr.Uid}
+		case tar.TypeReg:
 			buf := bytes.Buffer{}
 			io.Copy(&buf, tr)
-			elem.contents = buf.Bytes()
+			elem = file{name: hdr.Name, uid: hdr.Uid, contents: buf.Bytes()}
 		}
 		ret = append(ret, elem)
 	}
 	return ret
 }
-
-func TestRootFS(t *testing.T) {
-	layer0 := tarball{
-		dir{name: "/", uid: 0},
-		file{name: "/file", uid: 1, contents: []byte("from 0")},
-	}
-
-	layer1 := tarball{
-		file{name: "/file", uid: 0, contents: []byte("from 1")},
-	}
-
-	layer2 := tarball{
-		dir{name: "/", uid: 1},
-	}
-
-	tests := []struct {
-		name    string
-		image   tarball
-		want    []file
-		wantErr string
-	}{
-		{
-			name:  "empty tarball",
-			image: tarball{manifest{}},
-			want:  []file{},
-		},
-		{
-			name:    "missing layer",
-			image:   tarball{manifest{"layer0/layer.tar"}},
-			wantErr: "bad or missing manifest.json",
-		},
-		{
-			name: "basic file overwrite, layer order mixed",
-			image: tarball{
-				file{name: "layer1/layer.tar", contents: layer1.bytes(t)},
-				file{name: "layer0/layer.tar", contents: layer0.bytes(t)},
-				manifest{"layer0/layer.tar", "layer1/layer.tar"},
-			},
-			want: []file{
-				{name: "/", uid: 0},
-				{name: "/file", uid: 0, contents: []byte("from 1")},
-			},
-		},
-		{
-			name: "directory overwrite retains original dir",
-			image: tarball{
-				file{name: "layer0/layer.tar", contents: layer0.bytes(t)},
-				file{name: "layer1/layer.tar", contents: layer1.bytes(t)},
-				file{name: "layer2/layer.tar", contents: layer2.bytes(t)},
-				manifest{"layer0/layer.tar", "layer1/layer.tar", "layer2/layer.tar"},
-			},
-			want: []file{
-				{name: "/", uid: 0},
-				{name: "/file", uid: 0, contents: []byte("from 1")},
-				{name: "/", uid: 1},
-			},
-		},
-	}
-
-	for _, tt := range tests {
-		t.Run(tt.name, func(t *testing.T) {
-			in := bytes.NewReader(tt.image.bytes(t))
-			out := bytes.Buffer{}
-
-			err := RootFS(in, &out)
-			if tt.wantErr != "" {
-				assert.EqualError(t, err, tt.wantErr)
-				return
-			}
-			require.NoError(t, err)
-			got := extract(t, &out)
-			assert.Equal(t, got, tt.want)
-		})
-	}
-}