small improvements to testing tools, specially InstallAgentForPolicy (#…

…3265)

* remove missing verb on printf-like call upgrade integration test testStandaloneUpgrade.
 * make InstallAgentForPolicy receive a context and use either the context timeout or a 5-minutes default.
 * add new pkg/testing/tools/textcontext package
        The new package is for helper functions dealing with contexts and testing.T.
        It provides a WithDeadline function which respects t.Deadline()

(cherry picked from commit f2eb134)

pkg/testing/tools/testcontext/testcontext.go

@@ -0,0 +1,24 @@
+// Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+// or more contributor license agreements. Licensed under the Elastic License;
+// you may not use this file except in compliance with the Elastic License.
+
+package testcontext
+
+import (
+	"context"
+	"testing"
+	"time"
+)
+
+// WithDeadline returns a context with a deadline. The deadline is the earliest
+// of either the provided 'deadline' or t.Deadline().
+func WithDeadline(
+	t *testing.T,
+	parent context.Context,
+	deadline time.Time) (context.Context, context.CancelFunc) {
+	if d, ok := t.Deadline(); ok {
+		deadline = d
+	}
+	ctx, cancel := context.WithDeadline(parent, deadline)
+	return ctx, cancel
+}

pkg/testing/tools/tools.go

@@ -81,14 +81,21 @@ func InstallAgentWithPolicy(t *testing.T, ctx context.Context, installOpts atest
 		agentFixture.SetUninstallToken(uninstallToken)
 	}
 
-	err = InstallAgentForPolicy(t, installOpts, agentFixture, kibClient, policy.ID)
+	err = InstallAgentForPolicy(t, ctx, installOpts, agentFixture, kibClient, policy.ID)
 	return policy, err
 }
 
-// InstallAgentForPolicy enrolls the given agent
-// fixture in Fleet using the default Fleet Server, waits for the agent to be
-// online, and returns error or nil.
-func InstallAgentForPolicy(t *testing.T, installOpts atesting.InstallOpts, agentFixture *atesting.Fixture, kibClient *kibana.Client, policyID string) error {
+// InstallAgentForPolicy enrolls the provided agent fixture in Fleet using the
+// default Fleet Server, waits for the agent to come online, and returns either
+// an error or nil.
+// If the context (ctx) has a deadline, it will wait for the agent to become
+// online until the deadline of the context, or if not, a default 5-minute
+// deadline will be applied.
+func InstallAgentForPolicy(t *testing.T, ctx context.Context,
+	installOpts atesting.InstallOpts,
+	agentFixture *atesting.Fixture,
+	kibClient *kibana.Client,
+	policyID string) error {
 	t.Helper()
 
 	// Create enrollment API key
@@ -97,7 +104,7 @@ func InstallAgentForPolicy(t *testing.T, installOpts atesting.InstallOpts, agent
 	}
 
 	t.Logf("Creating enrollment API key...")
-	enrollmentToken, err := kibClient.CreateEnrollmentAPIKey(context.Background(), createEnrollmentAPIKeyReq)
+	enrollmentToken, err := kibClient.CreateEnrollmentAPIKey(ctx, createEnrollmentAPIKeyReq)
 	if err != nil {
 		return fmt.Errorf("unable to create enrollment API key: %w", err)
 	}
@@ -121,11 +128,15 @@ func InstallAgentForPolicy(t *testing.T, installOpts atesting.InstallOpts, agent
 	}
 	t.Logf(">>> Ran Enroll. Output: %s", output)
 
+	timeout := 5 * time.Minute
+	if deadline, ok := ctx.Deadline(); ok {
+		timeout = time.Until(deadline)
+	}
 	// Wait for Agent to be healthy
 	require.Eventually(
 		t,
 		WaitForAgentStatus(t, kibClient, "online"),
-		2*time.Minute,
+		timeout,
 		10*time.Second,
 		"Elastic Agent status is not online",
 	)

testing/integration/endpoint_security_test.go

@@ -30,6 +30,7 @@ import (
 	atesting "github.com/elastic/elastic-agent/pkg/testing"
 	"github.com/elastic/elastic-agent/pkg/testing/define"
 	"github.com/elastic/elastic-agent/pkg/testing/tools"
+	"github.com/elastic/elastic-agent/pkg/testing/tools/testcontext"
 )
 
 const (
@@ -77,7 +78,7 @@ func TestInstallAndCLIUninstallWithEndpointSecurity(t *testing.T) {
 		Isolate: false,
 		Sudo:    true, // requires Agent installation
 		OS: []define.OS{
-			define.OS{Type: define.Linux},
+			{Type: define.Linux},
 		},
 	})
 
@@ -101,7 +102,10 @@ func buildPolicyWithTamperProtection(policy kibana.AgentPolicy, protected bool)
 }
 
 func testInstallAndCLIUninstallWithEndpointSecurity(t *testing.T, info *define.Info, protected bool) {
-	t.Helper()
+	deadline := time.Now().Add(10 * time.Minute)
+	ctx, cancel := testcontext.WithDeadline(t, context.Background(), deadline)
+	defer cancel()
+
 	// Get path to agent executable.
 	fixture, err := define.NewFixture(t, define.Version())
 	require.NoError(t, err)
@@ -127,19 +131,16 @@ func testInstallAndCLIUninstallWithEndpointSecurity(t *testing.T, info *define.I
 		Force:          true,
 	}
 
-	ctx, cn := context.WithCancel(context.Background())
-	defer cn()
-
-	// Create policy
-	policy, err := tools.InstallAgentWithPolicy(t, ctx, installOpts, fixture, info.KibanaClient, createPolicyReq)
-	require.NoError(t, err)
+	policy, err := tools.InstallAgentWithPolicy(t, ctx,
+		installOpts, fixture, info.KibanaClient, createPolicyReq)
+	require.NoError(t, err, "failed to install agent with policy")
 
 	t.Log("Installing Elastic Defend")
 	pkgPolicyResp, err := installElasticDefendPackage(t, info, policy.ID)
 	require.NoErrorf(t, err, "Policy Response was: %v", pkgPolicyResp)
 
 	t.Log("Polling for endpoint-security to become Healthy")
-	ctx, cancel := context.WithTimeout(context.Background(), endpointHealthPollingTimeout)
+	ctx, cancel = context.WithTimeout(ctx, endpointHealthPollingTimeout)
 	defer cancel()
 
 	agentClient := fixture.Client()

testing/integration/install_test.go

@@ -53,8 +53,8 @@ func TestInstallWithoutBasePath(t *testing.T) {
 	}
 
 	topPath := filepath.Join(defaultBasePath, "Elastic", "Agent")
-	_, err = os.Stat(topPath)
-	require.True(t, os.IsNotExist(err))
+	err = os.RemoveAll(topPath)
+	require.NoError(t, err, "failed to remove %q. The test requires this path not to exist.")
 
 	// Run `elastic-agent install`.  We use `--force` to prevent interactive
 	// execution.

testing/integration/upgrade_test.go

@@ -466,7 +466,8 @@ func testStandaloneUpgrade(
 	}
 
 	upgradeTriggerOutput, err := f.Exec(ctx, upgradeCmdArgs)
-	require.NoErrorf(t, err, "error triggering agent upgrade to version %q, output:\n%s%", parsedUpgradeVersion, upgradeTriggerOutput)
+	require.NoErrorf(t, err, "error triggering agent upgrade to version %q, output:\n%s",
+		parsedUpgradeVersion, upgradeTriggerOutput)
 
 	require.Eventuallyf(t, func() bool {
 		return checkAgentHealthAndVersion(t, ctx, f, parsedUpgradeVersion.CoreVersion(), parsedUpgradeVersion.IsSnapshot(), expectedAgentHashAfterUpgrade)