1 files changed, 46 insertions, 0 deletions
diff --git a/workhorse/internal/upload/destination/reader_test.go b/workhorse/internal/upload/destination/reader_test.go
new file mode 100644
index 00000000000..a26f7746a13
--- /dev/null
+++ b/workhorse/internal/upload/destination/reader_test.go
@@ -0,0 +1,46 @@
+package destination
+
+import (
+	"fmt"
+	"io/ioutil"
+	"strings"
+	"testing"
+	"testing/iotest"
+
+	"github.com/stretchr/testify/require"
+)
+
+func TestHardLimitReader(t *testing.T) {
+	const text = "hello world"
+	r := iotest.OneByteReader(
+		&hardLimitReader{
+			r: strings.NewReader(text),
+			n: int64(len(text)),
+		},
+	)
+
+	out, err := ioutil.ReadAll(r)
+	require.NoError(t, err)
+	require.Equal(t, text, string(out))
+}
+
+func TestHardLimitReaderFail(t *testing.T) {
+	const text = "hello world"
+
+	for bufSize := len(text) / 2; bufSize < len(text)*2; bufSize++ {
+		t.Run(fmt.Sprintf("bufsize:%d", bufSize), func(t *testing.T) {
+			r := &hardLimitReader{
+				r: iotest.DataErrReader(strings.NewReader(text)),
+				n: int64(len(text)) - 1,
+			}
+			buf := make([]byte, bufSize)
+
+			var err error
+			for i := 0; err == nil && i < 1000; i++ {
+				_, err = r.Read(buf)
+			}
+
+			require.Equal(t, ErrEntityTooLarge, err)
+		})
+	}
+}