From 02d5fc5fd0de7934243f155c016a046f827b0fad Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Miloslav=20Trma=C4=8D?= <mitr@redhat.com>
Date: Mon, 13 Feb 2023 22:24:04 +0100
Subject: [PATCH] Add unit tests for the retry heuristics
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Signed-off-by: Miloslav Trmač <mitr@redhat.com>
---
 docker/body_reader_test.go | 86 ++++++++++++++++++++++++++++++++++++++
 1 file changed, 86 insertions(+)

diff --git a/docker/body_reader_test.go b/docker/body_reader_test.go
index ef702edc7..0011582b7 100644
--- a/docker/body_reader_test.go
+++ b/docker/body_reader_test.go
@@ -1,9 +1,13 @@
 package docker
 
 import (
+	"errors"
+	"math"
 	"net/http"
 	"testing"
+	"time"
 
+	"github.com/sirupsen/logrus"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 )
@@ -108,3 +112,85 @@ func TestParseContentRange(t *testing.T) {
 		assert.Error(t, err, c, c)
 	}
 }
+
+func TestMillisecondsSinceOptional(t *testing.T) {
+	current := time.Date(2023, 2, 9, 8, 7, 6, 5, time.UTC)
+	res := millisecondsSinceOptional(current, time.Time{})
+	assert.True(t, math.IsNaN(res))
+	tm := current.Add(-60 * time.Second) // 60 seconds _before_ current
+	res = millisecondsSinceOptional(current, tm)
+	assert.Equal(t, res, 60_000.0)
+}
+
+func TestBodyReaderErrorIfNotReconnecting(t *testing.T) {
+	// Silence logrus.Info logs in the tested method
+	prevLevel := logrus.StandardLogger().Level
+	logrus.StandardLogger().SetLevel(logrus.WarnLevel)
+	t.Cleanup(func() {
+		logrus.StandardLogger().SetLevel(prevLevel)
+	})
+
+	for _, c := range []struct {
+		name            string
+		previousRetry   bool
+		currentOffset   int64
+		currentTime     int // milliseconds
+		expectReconnect bool
+	}{
+		{
+			name:            "A lot of progress, after a long time, second retry",
+			previousRetry:   true,
+			currentOffset:   2 * bodyReaderMinimumProgress,
+			currentTime:     2 * bodyReaderMSSinceLastRetry,
+			expectReconnect: true,
+		},
+		{
+			name:            "A lot of progress, after little time, second retry",
+			previousRetry:   true,
+			currentOffset:   2 * bodyReaderMinimumProgress,
+			currentTime:     1,
+			expectReconnect: true,
+		},
+		{
+			name:            "Little progress, after a long time, second retry",
+			previousRetry:   true,
+			currentOffset:   1,
+			currentTime:     2 * bodyReaderMSSinceLastRetry,
+			expectReconnect: true,
+		},
+		{
+			name:            "Little progress, after little time, second retry",
+			previousRetry:   true,
+			currentOffset:   1,
+			currentTime:     1,
+			expectReconnect: false,
+		},
+		{
+			name:            "Little progress, after little time, first retry",
+			previousRetry:   false,
+			currentOffset:   1,
+			currentTime:     bodyReaderMSSinceLastRetry / 2,
+			expectReconnect: true,
+		},
+	} {
+		tm := time.Now()
+		br := bodyReader{}
+		if c.previousRetry {
+			br.lastRetryOffset = 2 * bodyReaderMinimumProgress
+			br.offset = br.lastRetryOffset + c.currentOffset
+			br.firstConnectionTime = tm.Add(-time.Duration(c.currentTime+2*bodyReaderMSSinceLastRetry) * time.Millisecond)
+			br.lastRetryTime = tm.Add(-time.Duration(c.currentTime) * time.Millisecond)
+		} else {
+			br.lastRetryOffset = -1
+			br.lastRetryTime = time.Time{}
+			br.offset = c.currentOffset
+			br.firstConnectionTime = tm.Add(-time.Duration(c.currentTime) * time.Millisecond)
+		}
+		err := br.errorIfNotReconnecting(errors.New("some error for error text only"), "URL for error text only")
+		if c.expectReconnect {
+			assert.NoError(t, err, c.name, br)
+		} else {
+			assert.Error(t, err, c.name, br)
+		}
+	}
+}