From cd67a98ec7dff231e956ca7ebb731b84da632cea Mon Sep 17 00:00:00 2001
From: Derek McGowan <derek@mcg.dev>
Date: Fri, 19 Aug 2022 15:59:25 -0700
Subject: [PATCH] Add transfer integration test

Signed-off-by: Derek McGowan <derek@mcg.dev>
---
 integration/client/transfer_test.go | 95 +++++++++++++++++++++++++++++
 1 file changed, 95 insertions(+)
 create mode 100644 integration/client/transfer_test.go

diff --git a/integration/client/transfer_test.go b/integration/client/transfer_test.go
new file mode 100644
index 000000000..74d19eeaa
--- /dev/null
+++ b/integration/client/transfer_test.go
@@ -0,0 +1,95 @@
+/*
+   Copyright The containerd Authors.
+
+   Licensed under the Apache License, Version 2.0 (the "License");
+   you may not use this file except in compliance with the License.
+   You may obtain a copy of the License at
+
+       http://www.apache.org/licenses/LICENSE-2.0
+
+   Unless required by applicable law or agreed to in writing, software
+   distributed under the License is distributed on an "AS IS" BASIS,
+   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+   See the License for the specific language governing permissions and
+   limitations under the License.
+*/
+
+package client
+
+import (
+	"bytes"
+	"context"
+	"fmt"
+	"io"
+	"testing"
+
+	"github.com/containerd/containerd"
+	"github.com/containerd/containerd/pkg/transfer/archive"
+)
+
+func TestTransferEcho(t *testing.T) {
+	client, err := newClient(t, address)
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer client.Close()
+
+	ctx, cancel := testContext(t)
+	defer cancel()
+	t.Run("ImportExportEchoBig", newImportExportEcho(ctx, client, bytes.Repeat([]byte("somecontent"), 17*1024)))
+	t.Run("ImportExportEchoSmall", newImportExportEcho(ctx, client, []byte("somecontent")))
+	t.Run("ImportExportEchoEmpty", newImportExportEcho(ctx, client, []byte("")))
+
+}
+
+func newImportExportEcho(ctx context.Context, client *containerd.Client, expected []byte) func(*testing.T) {
+	return func(t *testing.T) {
+		testBuf := newWaitBuffer()
+		err := client.Transfer(ctx, archive.NewImageImportStream(bytes.NewReader(expected)), archive.NewImageExportStream(testBuf))
+		if err != nil {
+			t.Fatal(err)
+		}
+		if !bytes.Equal(expected, testBuf.Bytes()) {
+			t.Fatalf("Bytes did not match\n\tActual:   %v\n\tExpected: %v", displayBytes(testBuf.Bytes()), displayBytes(expected))
+		}
+	}
+}
+
+type WriteBytesCloser interface {
+	io.WriteCloser
+	Bytes() []byte
+}
+
+func newWaitBuffer() WriteBytesCloser {
+	return &waitBuffer{
+		Buffer: bytes.NewBuffer(nil),
+		closed: make(chan struct{}),
+	}
+}
+
+type waitBuffer struct {
+	*bytes.Buffer
+	closed chan struct{}
+}
+
+func (wb *waitBuffer) Close() error {
+	select {
+	case <-wb.closed:
+	default:
+		close(wb.closed)
+	}
+	return nil
+}
+
+func (wb *waitBuffer) Bytes() []byte {
+	<-wb.closed
+	return wb.Buffer.Bytes()
+}
+
+func displayBytes(b []byte) string {
+	if len(b) > 40 {
+		skipped := len(b) - 20
+		return fmt.Sprintf("%s...(skipped %d)...%s", b[:10], skipped, b[len(b)-10:])
+	}
+	return string(b)
+}