From fafda48905d0a1cbcef19ffcf0c2e2f1679546e2 Mon Sep 17 00:00:00 2001
From: Phill MV <phillmv@github.com>
Date: Mon, 21 Oct 2024 12:32:57 -0400
Subject: [PATCH] added test for verifying we do 3 retries when fetching
 attestations.

---
 pkg/cmd/attestation/api/client_test.go        | 39 +++++++++++++++++--
 .../attestation/api/mock_apiClient_test.go    | 10 +++++
 2 files changed, 46 insertions(+), 3 deletions(-)

diff --git a/pkg/cmd/attestation/api/client_test.go b/pkg/cmd/attestation/api/client_test.go
index 4751a3ea8..e6bca4406 100644
--- a/pkg/cmd/attestation/api/client_test.go
+++ b/pkg/cmd/attestation/api/client_test.go
@@ -212,13 +212,12 @@ func TestGetAttestationsRetries(t *testing.T) {
 	fetcher := mockDataGenerator{
 		NumAttestations: 5,
 	}
-	l := io.NewTestHandler()
 
 	c := &LiveClient{
 		api: mockAPIClient{
 			OnRESTWithNext: fetcher.FlakyOnRESTSuccessWithNextPageHandler(),
 		},
-		logger: l,
+		logger: io.NewTestHandler(),
 	}
 
 	attestations, err := c.GetByRepoAndDigest(testRepo, testDigest, DefaultLimit)
@@ -229,9 +228,43 @@ func TestGetAttestationsRetries(t *testing.T) {
 	fetcher.AssertNumberOfCalls(t, "FlakyOnRESTSuccessWithNextPage:error", 2)
 
 	// but we still successfully got the right data
-	require.Equal(t, 10, len(attestations))
+	require.Equal(t, len(attestations), 10)
 	bundle := (attestations)[0].Bundle
 	require.Equal(t, bundle.GetMediaType(), "application/vnd.dev.sigstore.bundle.v0.3+json")
 
+	// same test as above, but for GetByOwnerAndDigest:
+	attestations, err = c.GetByOwnerAndDigest(testOwner, testDigest, DefaultLimit)
+	require.NoError(t, err)
+
+	// because we haven't reset the mock, we have added 2 more failed requests
+	fetcher.AssertNumberOfCalls(t, "FlakyOnRESTSuccessWithNextPage:error", 4)
+
+	require.Equal(t, len(attestations), 10)
+	bundle = (attestations)[0].Bundle
+	require.Equal(t, bundle.GetMediaType(), "application/vnd.dev.sigstore.bundle.v0.3+json")
+
+	getAttestationRetryInterval = oldInterval
+}
+
+// test total retries
+func TestGetAttestationsMaxRetries(t *testing.T) {
+	oldInterval := getAttestationRetryInterval
+	getAttestationRetryInterval = 0
+
+	fetcher := mockDataGenerator{
+		NumAttestations: 5,
+	}
+
+	c := &LiveClient{
+		api: mockAPIClient{
+			OnRESTWithNext: fetcher.OnREST500ErrorHandler(),
+		},
+		logger: io.NewTestHandler(),
+	}
+
+	_, err := c.GetByRepoAndDigest(testRepo, testDigest, DefaultLimit)
+	require.Error(t, err)
+
+	fetcher.AssertNumberOfCalls(t, "OnREST500Error", 4)
 	getAttestationRetryInterval = oldInterval
 }
diff --git a/pkg/cmd/attestation/api/mock_apiClient_test.go b/pkg/cmd/attestation/api/mock_apiClient_test.go
index e2364f048..0add67b75 100644
--- a/pkg/cmd/attestation/api/mock_apiClient_test.go
+++ b/pkg/cmd/attestation/api/mock_apiClient_test.go
@@ -65,6 +65,16 @@ func (m *mockDataGenerator) FlakyOnRESTSuccessWithNextPageHandler() func(hostnam
 	}
 }
 
+// always returns a 500
+func (m *mockDataGenerator) OnREST500ErrorHandler() func(hostname, method, p string, body io.Reader, data interface{}) (string, error) {
+	m.On("OnREST500Error").Return()
+	return func(hostname, method, p string, body io.Reader, data interface{}) (string, error) {
+		m.MethodCalled("OnREST500Error")
+
+		return "", cliAPI.HTTPError{HTTPError: &ghAPI.HTTPError{StatusCode: 500}}
+	}
+}
+
 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++ {