From 664e09fdbcccedc50916ca533acb700e39e7fa9b Mon Sep 17 00:00:00 2001
From: Phill MV <phillmv@github.com>
Date: Mon, 21 Oct 2024 11:20:46 -0400
Subject: [PATCH] wip: added test that fails in the absence of a backoff.

---
 pkg/cmd/attestation/api/client_test.go        | 26 +++++++++++++++++++
 .../attestation/api/mock_apiClient_test.go    | 25 ++++++++++++++++++
 2 files changed, 51 insertions(+)

diff --git a/pkg/cmd/attestation/api/client_test.go b/pkg/cmd/attestation/api/client_test.go
index bfcb40f5a..5544f7476 100644
--- a/pkg/cmd/attestation/api/client_test.go
+++ b/pkg/cmd/attestation/api/client_test.go
@@ -204,3 +204,29 @@ func TestGetTrustDomain(t *testing.T) {
 	})
 
 }
+
+func TestGetAttestationsRetries(t *testing.T) {
+	fetcher := mockDataGenerator{
+		NumAttestations: 5,
+	}
+	l := io.NewTestHandler()
+
+	c := &LiveClient{
+		api: mockAPIClient{
+			OnRESTWithNext: fetcher.FlakyOnRESTSuccessWithNextPageHandler(),
+		},
+		logger: l,
+	}
+
+	attestations, err := c.GetByRepoAndDigest(testRepo, testDigest, DefaultLimit)
+	require.NoError(t, err)
+
+	// assert the error path was executed; because this is a paged
+	// request, it should have errored twice
+	fetcher.AssertNumberOfCalls(t, "FlakyOnRESTSuccessWithNextPage:error", 2)
+
+	// but we still successfully got the right data
+	require.Equal(t, 10, len(attestations))
+	bundle := (attestations)[0].Bundle
+	require.Equal(t, bundle.GetMediaType(), "application/vnd.dev.sigstore.bundle.v0.3+json")
+}
diff --git a/pkg/cmd/attestation/api/mock_apiClient_test.go b/pkg/cmd/attestation/api/mock_apiClient_test.go
index d58cdbc79..e2364f048 100644
--- a/pkg/cmd/attestation/api/mock_apiClient_test.go
+++ b/pkg/cmd/attestation/api/mock_apiClient_test.go
@@ -6,6 +6,10 @@ import (
 	"fmt"
 	"io"
 	"strings"
+
+	cliAPI "github.com/cli/cli/v2/api"
+	ghAPI "github.com/cli/go-gh/v2/pkg/api"
+	"github.com/stretchr/testify/mock"
 )
 
 type mockAPIClient struct {
@@ -22,6 +26,7 @@ func (m mockAPIClient) REST(hostname, method, p string, body io.Reader, data int
 }
 
 type mockDataGenerator struct {
+	mock.Mock
 	NumAttestations int
 }
 
@@ -40,6 +45,26 @@ func (m mockDataGenerator) OnRESTSuccessWithNextPage(hostname, method, p string,
 	return m.OnRESTWithNextSuccessHelper(hostname, method, p, body, data, false)
 }
 
+// Returns a func that just calls OnRESTSuccessWithNextPage but half the time
+// it returns a 500 error.
+func (m *mockDataGenerator) FlakyOnRESTSuccessWithNextPageHandler() func(hostname, method, p string, body io.Reader, data interface{}) (string, error) {
+	// set up the flake counter
+	m.On("FlakyOnRESTSuccessWithNextPage:error").Return()
+
+	count := 0
+	return func(hostname, method, p string, body io.Reader, data interface{}) (string, error) {
+		if count%2 == 0 {
+			m.MethodCalled("FlakyOnRESTSuccessWithNextPage:error")
+
+			count = count + 1
+			return "", cliAPI.HTTPError{HTTPError: &ghAPI.HTTPError{StatusCode: 500}}
+		} else {
+			count = count + 1
+			return m.OnRESTSuccessWithNextPage(hostname, method, p, body, data)
+		}
+	}
+}
+
 func (m mockDataGenerator) OnRESTWithNextSuccessHelper(hostname, method, p string, body io.Reader, data interface{}, hasNext bool) (string, error) {
 	atts := make([]*Attestation, m.NumAttestations)
 	for j := 0; j < m.NumAttestations; j++ {