From 6e4b379cdb6bdb852dfd910647cf87c71bcb4b0f Mon Sep 17 00:00:00 2001
From: Anderson Queiroz <anderson.queiroz@elastic.co>
Date: Wed, 27 Mar 2024 14:34:46 +0100
Subject: [PATCH] add TEST_PACKAGES environment variable for integration tests
 (#4480)

---
 magefile.go                  | 18 ++++++++++++++++++
 pkg/testing/runner/config.go |  5 +++++
 pkg/testing/runner/runner.go | 22 +++++++++++++++++++---
 3 files changed, 42 insertions(+), 3 deletions(-)

diff --git a/magefile.go b/magefile.go
index 3d15a6771f3..14b44e20737 100644
--- a/magefile.go
+++ b/magefile.go
@@ -2285,6 +2285,7 @@ func createTestRunner(matrix bool, singleTest string, goTestFlags string, batche
 		DiagnosticsDir: diagDir,
 		StateDir:       ".integration-cache",
 		Platforms:      testPlatforms(),
+		Packages:       testPackages(),
 		Groups:         testGroups(),
 		Matrix:         matrix,
 		SingleTest:     singleTest,
@@ -2380,6 +2381,23 @@ func testPlatforms() []string {
 	return platforms
 }
 
+func testPackages() []string {
+	packagesStr, defined := os.LookupEnv("TEST_PACKAGES")
+	if !defined {
+		return nil
+	}
+
+	var packages []string
+	for _, p := range strings.Split(packagesStr, ",") {
+		if p == "tar.gz" {
+			p = "targz"
+		}
+		packages = append(packages, p)
+	}
+
+	return packages
+}
+
 func testGroups() []string {
 	groupsStr := os.Getenv("TEST_GROUPS")
 	if groupsStr == "" {
diff --git a/pkg/testing/runner/config.go b/pkg/testing/runner/config.go
index 22494ed58e3..92229e35a99 100644
--- a/pkg/testing/runner/config.go
+++ b/pkg/testing/runner/config.go
@@ -28,6 +28,11 @@ type Config struct {
 	// defined in this list.
 	Platforms []string
 
+	// Packages filters the tests to only run on the provided list
+	// of platforms even if the tests supports more than what is
+	// defined in this list.
+	Packages []string
+
 	// BinaryName is the name of the binary package under test, i.e, elastic-agent, metricbeat, etc
 	// this is used to copy the .tar.gz to the remote host
 	BinaryName string
diff --git a/pkg/testing/runner/runner.go b/pkg/testing/runner/runner.go
index 939cf108e46..b906e3496a5 100644
--- a/pkg/testing/runner/runner.go
+++ b/pkg/testing/runner/runner.go
@@ -447,20 +447,36 @@ func (r *Runner) validate() error {
 
 // getBuilds returns the build for the batch.
 func (r *Runner) getBuilds(b OSBatch) []Build {
-	builds := []Build{}
+	var builds []Build
 	formats := []string{"targz", "zip", "rpm", "deb"}
 	binaryName := "elastic-agent"
 
+	var packages []string
+	for _, p := range r.cfg.Packages {
+		if slices.Contains(formats, p) {
+			packages = append(packages, p)
+		}
+	}
+	if len(packages) == 0 {
+		packages = formats
+	}
+
 	// This is for testing beats in serverless environment
 	if strings.HasSuffix(r.cfg.BinaryName, "beat") {
-		formats = []string{"targz", "zip"}
+		var serverlessPackages []string
+		for _, p := range packages {
+			if slices.Contains([]string{"targz", "zip"}, p) {
+				packages = append(packages, p)
+			}
+		}
+		packages = serverlessPackages
 	}
 
 	if r.cfg.BinaryName != "" {
 		binaryName = r.cfg.BinaryName
 	}
 
-	for _, f := range formats {
+	for _, f := range packages {
 		arch := b.OS.Arch
 		if arch == define.AMD64 {
 			arch = "x86_64"