Integration Testing Framework: Add retries for artifacts API call (#4348)

* Add retries for artifacts API call

* Add newline

* Adding logging

* Pass logger in constructor

* Add missing arg

* Fix logic for happy path

* Adding comment

* Do not keep retrying if request context was cancelled or timed out

Author: ycombinator

pkg/testing/tools/artifacts_api.go

@@ -14,6 +14,7 @@ import (
 	"net/url"
 	"runtime"
 	"sort"
+	"time"
 
 	"github.com/elastic/elastic-agent/pkg/testing"
 	"github.com/elastic/elastic-agent/pkg/version"
@@ -29,6 +30,9 @@ const (
 
 	artifactElasticAgentProject = "elastic-agent-package"
 	artifactReleaseCDN          = "https://artifacts.elastic.co/downloads/beats/elastic-agent"
+
+	maxAttemptsForArtifactsAPICall   = 6
+	retryIntervalForArtifactsAPICall = 5 * time.Second
 )
 
 var (
@@ -137,14 +141,17 @@ type ArtifactAPIClient struct {
 	c      httpDoer
 	url    string
 	cdnURL string
+
+	logger logger
 }
 
 // NewArtifactAPIClient creates a new Artifact API client
-func NewArtifactAPIClient(opts ...ArtifactAPIClientOpt) *ArtifactAPIClient {
+func NewArtifactAPIClient(logger logger, opts ...ArtifactAPIClientOpt) *ArtifactAPIClient {
 	c := &ArtifactAPIClient{
 		url:    defaultArtifactAPIURL,
 		cdnURL: artifactReleaseCDN,
 		c:      new(http.Client),
+		logger: logger,
 	}
 
 	for _, opt := range opts {
@@ -310,9 +317,37 @@ func (aac ArtifactAPIClient) createAndPerformRequest(ctx context.Context, URL st
 		return nil, err
 	}
 
-	resp, err := aac.c.Do(req)
+	// Make the request with retries as the artifacts API can sometimes be flaky.
+	var resp *http.Response
+	// TODO (once we're on Go 1.22): replace with for numAttempts := range maxAttemptsForArtifactsAPICall {
+	for numAttempts := 0; numAttempts < maxAttemptsForArtifactsAPICall; numAttempts++ {
+		resp, err = aac.c.Do(req)
+
+		// If there is no error, no need to retry the request.
+		if err == nil {
+			break
+		}
+
+		// If the context was cancelled or timed out, return early
+		if errors.Is(err, context.Canceled) || errors.Is(err, context.DeadlineExceeded) {
+			return nil, fmt.Errorf(
+				"executing http request %s %s (attempt %d of %d) cancelled or timed out: %w",
+				req.Method, req.URL, numAttempts+1, maxAttemptsForArtifactsAPICall, err,
+			)
+		}
+
+		aac.logger.Logf(
+			"failed attempt %d of %d executing http request %s %s: %s; retrying after %v...",
+			numAttempts+1, maxAttemptsForArtifactsAPICall, req.Method, req.URL, err.Error(), retryIntervalForArtifactsAPICall,
+		)
+		time.Sleep(retryIntervalForArtifactsAPICall)
+	}
+
 	if err != nil {
-		return nil, fmt.Errorf("executing http request %v: %w", req, err)
+		return nil, fmt.Errorf(
+			"failed executing http request %s %s after %d attempts: %w",
+			req.Method, req.URL, maxAttemptsForArtifactsAPICall, err,
+		)
 	}
 
 	return resp, nil
@@ -341,7 +376,7 @@ type logger interface {
 	Logf(format string, args ...any)
 }
 
-func (aac ArtifactAPIClient) GetLatestSnapshotVersion(ctx context.Context, log logger) (*version.ParsedSemVer, error) {
+func (aac ArtifactAPIClient) GetLatestSnapshotVersion(ctx context.Context) (*version.ParsedSemVer, error) {
 	vList, err := aac.GetVersions(ctx)
 	if err != nil {
 		return nil, err
@@ -355,7 +390,7 @@ func (aac ArtifactAPIClient) GetLatestSnapshotVersion(ctx context.Context, log l
 	for _, v := range vList.Versions {
 		pv, err := version.ParseVersion(v)
 		if err != nil {
-			log.Logf("invalid version retrieved from artifact API: %q", v)
+			aac.logger.Logf("invalid version retrieved from artifact API: %q", v)
 			return nil, ErrInvalidVersionRetrieved
 		}
 		sortedParsedVersions = append(sortedParsedVersions, pv)

pkg/testing/tools/artifacts_api_test.go

@@ -171,7 +171,7 @@ func TestDefaultArtifactAPIClientErrorHttpStatus(t *testing.T) {
 			testSrv := httptest.NewServer(errorHandler)
 			defer testSrv.Close()
 
-			aac := NewArtifactAPIClient(WithUrl(testSrv.URL))
+			aac := NewArtifactAPIClient(t, WithUrl(testSrv.URL))
 			_, err := aac.GetVersions(context.Background())
 			assert.ErrorIs(t, err, ErrBadHTTPStatusCode, "Expected ErrBadHTTPStatusCode for status code %d", httpErrorCode)
 			_, err = aac.GetBuildsForVersion(context.Background(), "1.2.3-SNAPSHOT")
@@ -201,7 +201,7 @@ func TestDefaultArtifactAPIClient(t *testing.T) {
 	testSrv := httptest.NewServer(cannedRespHandler)
 	defer testSrv.Close()
 
-	aac := NewArtifactAPIClient(WithUrl(testSrv.URL))
+	aac := NewArtifactAPIClient(t, WithUrl(testSrv.URL))
 	versions, err := aac.GetVersions(context.Background())
 	assert.NoError(t, err)
 	assert.NotNil(t, versions)
@@ -245,7 +245,7 @@ func TestRemoveUnreleasedVersions(t *testing.T) {
 		resp.WriteHeader(http.StatusOK)
 	}))
 
-	client := NewArtifactAPIClient(WithCDNUrl(server.URL))
+	client := NewArtifactAPIClient(t, WithCDNUrl(server.URL))
 
 	t.Run("removes unreleased versions", func(t *testing.T) {
 		unreleasedVersions = map[string]struct{}{

testing/integration/upgrade_broken_package_test.go

@@ -41,7 +41,7 @@ func TestUpgradeBrokenPackageVersion(t *testing.T) {
 	require.NoError(t, err)
 
 	// Upgrade to an old build.
-	upgradeToVersion, err := upgradetest.PreviousMinor(ctx, define.Version())
+	upgradeToVersion, err := upgradetest.PreviousMinor(ctx, define.Version(), t)
 	require.NoError(t, err)
 	endFixture, err := atesting.NewFixture(
 		t,

testing/integration/upgrade_downgrade_test.go

@@ -42,8 +42,8 @@ func TestStandaloneDowngradeToSpecificSnapshotBuild(t *testing.T) {
 	ctx, cancel := testcontext.WithDeadline(t, context.Background(), time.Now().Add(10*time.Minute))
 	defer cancel()
 
-	aac := tools.NewArtifactAPIClient()
-	latestSnapshotVersion, err := aac.GetLatestSnapshotVersion(ctx, t)
+	aac := tools.NewArtifactAPIClient(t)
+	latestSnapshotVersion, err := aac.GetLatestSnapshotVersion(ctx)
 	require.NoError(t, err)
 
 	// start at the build version as we want to test the retry

testing/integration/upgrade_gpg_test.go

@@ -114,7 +114,7 @@ func TestStandaloneUpgradeWithGPGFallbackOneRemoteFailing(t *testing.T) {
 	require.NoError(t, err)
 
 	// Upgrade to an old build.
-	upgradeToVersion, err := upgradetest.PreviousMinor(ctx, define.Version())
+	upgradeToVersion, err := upgradetest.PreviousMinor(ctx, define.Version(), t)
 	require.NoError(t, err)
 	endFixture, err := atesting.NewFixture(
 		t,

testing/integration/upgrade_rollback_test.go

@@ -51,7 +51,7 @@ func TestStandaloneUpgradeRollback(t *testing.T) {
 
 	// Upgrade from an old build because the new watcher from the new build will
 	// be ran. Otherwise the test will run the old watcher from the old build.
-	upgradeFromVersion, err := upgradetest.PreviousMinor(ctx, define.Version())
+	upgradeFromVersion, err := upgradetest.PreviousMinor(ctx, define.Version(), t)
 	require.NoError(t, err)
 	startFixture, err := atesting.NewFixture(
 		t,
@@ -166,7 +166,7 @@ func TestStandaloneUpgradeRollbackOnRestarts(t *testing.T) {
 
 	// Upgrade from an old build because the new watcher from the new build will
 	// be ran. Otherwise the test will run the old watcher from the old build.
-	upgradeFromVersion, err := upgradetest.PreviousMinor(ctx, define.Version())
+	upgradeFromVersion, err := upgradetest.PreviousMinor(ctx, define.Version(), t)
 	require.NoError(t, err)
 	startFixture, err := atesting.NewFixture(
 		t,

testing/integration/upgrade_standalone_inprogress_test.go

@@ -39,7 +39,7 @@ func TestStandaloneUpgradeFailsWhenUpgradeIsInProgress(t *testing.T) {
 	// than the current version and upgrade to the current version. Then we attempt
 	// upgrading to the current version again, expecting Elastic Agent to disallow
 	// this second upgrade.
-	upgradeFromVersion, err := upgradetest.PreviousMinor(ctx, define.Version())
+	upgradeFromVersion, err := upgradetest.PreviousMinor(ctx, define.Version(), t)
 	require.NoError(t, err)
 	startFixture, err := atesting.NewFixture(
 		t,

testing/integration/upgrade_standalone_retry_test.go

@@ -49,7 +49,7 @@ func TestStandaloneUpgradeRetryDownload(t *testing.T) {
 	require.NoError(t, err, "failed to get end agent build version info")
 
 	// Upgrade to an older snapshot build of same version but with a different commit hash
-	aac := tools.NewArtifactAPIClient()
+	aac := tools.NewArtifactAPIClient(t)
 	buildInfo, err := aac.FindBuild(ctx, startVersion.VersionWithPrerelease(), startVersionInfo.Binary.Commit, 0)
 	if errors.Is(err, tools.ErrBuildNotFound) {
 		t.Skipf("there is no other build with a non-matching commit hash in the given version %s", define.Version())

testing/integration/upgrade_standalone_test.go

@@ -33,7 +33,7 @@ func TestStandaloneUpgrade(t *testing.T) {
 	// test 2 current 8.x version and 1 previous 7.x version
 	ctx, cancel := testcontext.WithDeadline(t, context.Background(), time.Now().Add(1*time.Minute))
 	defer cancel()
-	versionList, err := upgradetest.GetUpgradableVersions(ctx, define.Version(), 2, 1)
+	versionList, err := upgradetest.GetUpgradableVersions(ctx, define.Version(), 2, 1, t)
 	require.NoError(t, err)
 	endVersion, err := version.ParseVersion(define.Version())
 	require.NoError(t, err)

testing/integration/upgrade_uninstall_test.go

@@ -51,7 +51,7 @@ func TestStandaloneUpgradeUninstallKillWatcher(t *testing.T) {
 	// Start on a snapshot build, we want this test to upgrade to our
 	// build to ensure that the uninstall will kill the watcher.
 	// We need a snapshot with a non-matching commit hash to perform the upgrade
-	aac := tools.NewArtifactAPIClient()
+	aac := tools.NewArtifactAPIClient(t)
 	buildInfo, err := aac.FindBuild(ctx, endVersion.VersionWithPrerelease(), endVersionInfo.Binary.Commit, 0)
 	if errors.Is(err, tools.ErrBuildNotFound) {
 		t.Skipf("there is no other build with a non-matching commit hash in the given version %s", endVersion.VersionWithPrerelease())

testing/upgradetest/versions.go

@@ -10,6 +10,7 @@ import (
 	"fmt"
 	"sort"
 	"strings"
+	"testing"
 
 	"github.com/elastic/elastic-agent/pkg/testing/tools"
 	"github.com/elastic/elastic-agent/pkg/version"
@@ -37,8 +38,8 @@ var (
 )
 
 // GetUpgradableVersions returns the version that the upgradeToVersion can upgrade from.
-func GetUpgradableVersions(ctx context.Context, upgradeToVersion string, currentMajorVersions int, previousMajorVersions int) ([]*version.ParsedSemVer, error) {
-	aac := tools.NewArtifactAPIClient()
+func GetUpgradableVersions(ctx context.Context, upgradeToVersion string, currentMajorVersions int, previousMajorVersions int, t *testing.T) ([]*version.ParsedSemVer, error) {
+	aac := tools.NewArtifactAPIClient(t)
 	vList, err := aac.GetVersions(ctx)
 	if err != nil {
 		return nil, fmt.Errorf("error retrieving versions from Artifact API: %w", err)
@@ -134,8 +135,8 @@ func getUpgradableVersions(ctx context.Context, vList *tools.VersionList, upgrad
 // PreviousMinor gets the previous minor version of the provided version.
 //
 // This checks with the artifact API to ensure to only return version that have actual builds.
-func PreviousMinor(ctx context.Context, version string) (string, error) {
-	versions, err := GetUpgradableVersions(ctx, version, 1, 0)
+func PreviousMinor(ctx context.Context, version string, t *testing.T) (string, error) {
+	versions, err := GetUpgradableVersions(ctx, version, 1, 0, t)
 	if err != nil {
 		return "", err
 	}