1 files changed, 149 insertions, 0 deletions
diff --git a/workhorse/internal/upload/destination/objectstore/object_test.go b/workhorse/internal/upload/destination/objectstore/object_test.go
new file mode 100644
index 00000000000..24117891b6d
--- /dev/null
+++ b/workhorse/internal/upload/destination/objectstore/object_test.go
@@ -0,0 +1,149 @@
+package objectstore_test
+
+import (
+	"context"
+	"io"
+	"net/http"
+	"net/http/httptest"
+	"strings"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/require"
+
+	"gitlab.com/gitlab-org/gitlab/workhorse/internal/upload/destination/objectstore"
+	"gitlab.com/gitlab-org/gitlab/workhorse/internal/upload/destination/objectstore/test"
+)
+
+const testTimeout = 10 * time.Second
+
+type osFactory func() (*test.ObjectstoreStub, *httptest.Server)
+
+func testObjectUploadNoErrors(t *testing.T, startObjectStore osFactory, useDeleteURL bool, contentType string) {
+	osStub, ts := startObjectStore()
+	defer ts.Close()
+
+	objectURL := ts.URL + test.ObjectPath
+	var deleteURL string
+	if useDeleteURL {
+		deleteURL = objectURL
+	}
+
+	putHeaders := map[string]string{"Content-Type": contentType}
+
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+
+	deadline := time.Now().Add(testTimeout)
+	object, err := objectstore.NewObject(objectURL, deleteURL, putHeaders, test.ObjectSize)
+	require.NoError(t, err)
+
+	// copy data
+	n, err := object.Consume(ctx, strings.NewReader(test.ObjectContent), deadline)
+	require.NoError(t, err)
+	require.Equal(t, test.ObjectSize, n, "Uploaded file mismatch")
+
+	require.Equal(t, contentType, osStub.GetHeader(test.ObjectPath, "Content-Type"))
+
+	// Checking MD5 extraction
+	require.Equal(t, osStub.GetObjectMD5(test.ObjectPath), object.ETag())
+
+	// Checking cleanup
+	cancel()
+	require.Equal(t, 1, osStub.PutsCnt(), "Object hasn't been uploaded")
+
+	var expectedDeleteCnt int
+	if useDeleteURL {
+		expectedDeleteCnt = 1
+	}
+	require.Eventually(t, func() bool { return osStub.DeletesCnt() == expectedDeleteCnt }, time.Second, time.Millisecond)
+
+	if useDeleteURL {
+		require.Equal(t, 1, osStub.DeletesCnt(), "Object hasn't been deleted")
+	} else {
+		require.Equal(t, 0, osStub.DeletesCnt(), "Object has been deleted")
+	}
+}
+
+func TestObjectUpload(t *testing.T) {
+	t.Run("with delete URL", func(t *testing.T) {
+		testObjectUploadNoErrors(t, test.StartObjectStore, true, "application/octet-stream")
+	})
+	t.Run("without delete URL", func(t *testing.T) {
+		testObjectUploadNoErrors(t, test.StartObjectStore, false, "application/octet-stream")
+	})
+	t.Run("with custom content type", func(t *testing.T) {
+		testObjectUploadNoErrors(t, test.StartObjectStore, false, "image/jpeg")
+	})
+	t.Run("with upcase ETAG", func(t *testing.T) {
+		factory := func() (*test.ObjectstoreStub, *httptest.Server) {
+			md5s := map[string]string{
+				test.ObjectPath: strings.ToUpper(test.ObjectMD5),
+			}
+
+			return test.StartObjectStoreWithCustomMD5(md5s)
+		}
+
+		testObjectUploadNoErrors(t, factory, false, "application/octet-stream")
+	})
+}
+
+func TestObjectUpload404(t *testing.T) {
+	ts := httptest.NewServer(http.NotFoundHandler())
+	defer ts.Close()
+
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+
+	deadline := time.Now().Add(testTimeout)
+	objectURL := ts.URL + test.ObjectPath
+	object, err := objectstore.NewObject(objectURL, "", map[string]string{}, test.ObjectSize)
+	require.NoError(t, err)
+	_, err = object.Consume(ctx, strings.NewReader(test.ObjectContent), deadline)
+
+	require.Error(t, err)
+	_, isStatusCodeError := err.(objectstore.StatusCodeError)
+	require.True(t, isStatusCodeError, "Should fail with StatusCodeError")
+	require.Contains(t, err.Error(), "404")
+}
+
+type endlessReader struct{}
+
+func (e *endlessReader) Read(p []byte) (n int, err error) {
+	for i := 0; i < len(p); i++ {
+		p[i] = '*'
+	}
+
+	return len(p), nil
+}
+
+// TestObjectUploadBrokenConnection purpose is to ensure that errors caused by the upload destination get propagated back correctly.
+// This is important for troubleshooting in production.
+func TestObjectUploadBrokenConnection(t *testing.T) {
+	// This test server closes connection immediately
+	ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		hj, ok := w.(http.Hijacker)
+		if !ok {
+			require.FailNow(t, "webserver doesn't support hijacking")
+		}
+		conn, _, err := hj.Hijack()
+		if err != nil {
+			http.Error(w, err.Error(), http.StatusInternalServerError)
+			return
+		}
+		conn.Close()
+	}))
+	defer ts.Close()
+
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+
+	deadline := time.Now().Add(testTimeout)
+	objectURL := ts.URL + test.ObjectPath
+	object, err := objectstore.NewObject(objectURL, "", map[string]string{}, -1)
+	require.NoError(t, err)
+
+	_, copyErr := object.Consume(ctx, &endlessReader{}, deadline)
+	require.Error(t, copyErr)
+	require.NotEqual(t, io.ErrClosedPipe, copyErr, "We are shadowing the real error")
+}