From d3ba6388b3d10bb3d4a57602167fcad17c91dbe6 Mon Sep 17 00:00:00 2001
From: VihasMakwana <121151420+VihasMakwana@users.noreply.github.com>
Date: Tue, 27 Aug 2024 00:02:59 +0530
Subject: [PATCH] [testing/integration] Fix flaky test cases (#5359)

* fix: fix concurrency issue with the TestUpgrade* tests

* fix: flaky long running test cases

* chore: minor optimization
---
 .../handlers/handler_action_upgrade_test.go   | 22 +++++++++++++------
 .../agent_long_running_leak_test.go           |  2 +-
 2 files changed, 16 insertions(+), 8 deletions(-)

diff --git a/internal/pkg/agent/application/actions/handlers/handler_action_upgrade_test.go b/internal/pkg/agent/application/actions/handlers/handler_action_upgrade_test.go
index e331e477253..2f127c699f0 100644
--- a/internal/pkg/agent/application/actions/handlers/handler_action_upgrade_test.go
+++ b/internal/pkg/agent/application/actions/handlers/handler_action_upgrade_test.go
@@ -7,7 +7,6 @@ package handlers
 import (
 	"context"
 	"testing"
-	"time"
 
 	"github.com/stretchr/testify/require"
 
@@ -26,7 +25,8 @@ import (
 )
 
 type mockUpgradeManager struct {
-	msgChan chan string
+	msgChan       chan string
+	completedChan chan struct{}
 }
 
 func (u *mockUpgradeManager) Upgradeable() bool {
@@ -39,7 +39,7 @@ func (u *mockUpgradeManager) Reload(rawConfig *config.Config) error {
 
 func (u *mockUpgradeManager) Upgrade(ctx context.Context, version string, sourceURI string, action *fleetapi.ActionUpgrade, details *details.Details, skipVerifyOverride bool, skipDefaultPgp bool, pgpBytes ...string) (_ reexec.ShutdownCallbackFn, err error) {
 	select {
-	case <-time.After(2 * time.Second):
+	case <-u.completedChan:
 		u.msgChan <- "completed " + version
 		return nil, nil
 	case <-ctx.Done():
@@ -66,6 +66,7 @@ func TestUpgradeHandler(t *testing.T) {
 
 	agentInfo := &info.AgentInfo{}
 	msgChan := make(chan string)
+	completedChan := make(chan struct{})
 
 	// Create and start the coordinator
 	c := coordinator.New(
@@ -75,7 +76,7 @@ func TestUpgradeHandler(t *testing.T) {
 		agentInfo,
 		component.RuntimeSpecs{},
 		nil,
-		&mockUpgradeManager{msgChan: msgChan},
+		&mockUpgradeManager{msgChan: msgChan, completedChan: completedChan},
 		nil, nil, nil, nil, nil, false)
 	//nolint:errcheck // We don't need the termination state of the Coordinator
 	go c.Run(ctx)
@@ -85,6 +86,8 @@ func TestUpgradeHandler(t *testing.T) {
 		Version: "8.3.0", SourceURI: "http://localhost"}}
 	ack := noopacker.New()
 	err := u.Handle(ctx, &a, ack)
+	// indicate that upgrade is completed
+	close(completedChan)
 	require.NoError(t, err)
 	msg := <-msgChan
 	require.Equal(t, "completed 8.3.0", msg)
@@ -100,6 +103,7 @@ func TestUpgradeHandlerSameVersion(t *testing.T) {
 
 	agentInfo := &info.AgentInfo{}
 	msgChan := make(chan string)
+	completedChan := make(chan struct{})
 
 	// Create and start the Coordinator
 	c := coordinator.New(
@@ -109,7 +113,7 @@ func TestUpgradeHandlerSameVersion(t *testing.T) {
 		agentInfo,
 		component.RuntimeSpecs{},
 		nil,
-		&mockUpgradeManager{msgChan: msgChan},
+		&mockUpgradeManager{msgChan: msgChan, completedChan: completedChan},
 		nil, nil, nil, nil, nil, false)
 	//nolint:errcheck // We don't need the termination state of the Coordinator
 	go c.Run(ctx)
@@ -122,6 +126,8 @@ func TestUpgradeHandlerSameVersion(t *testing.T) {
 	err2 := u.Handle(ctx, &a, ack)
 	require.NoError(t, err1)
 	require.NoError(t, err2)
+	// indicate that upgrade is completed
+	close(completedChan)
 	msg := <-msgChan
 	require.Equal(t, "completed 8.3.0", msg)
 }
@@ -136,6 +142,7 @@ func TestUpgradeHandlerNewVersion(t *testing.T) {
 
 	agentInfo := &info.AgentInfo{}
 	msgChan := make(chan string)
+	completedChan := make(chan struct{})
 
 	// Create and start the Coordinator
 	c := coordinator.New(
@@ -145,7 +152,7 @@ func TestUpgradeHandlerNewVersion(t *testing.T) {
 		agentInfo,
 		component.RuntimeSpecs{},
 		nil,
-		&mockUpgradeManager{msgChan: msgChan},
+		&mockUpgradeManager{msgChan: msgChan, completedChan: completedChan},
 		nil, nil, nil, nil, nil, false)
 	//nolint:errcheck // We don't need the termination state of the Coordinator
 	go c.Run(ctx)
@@ -158,11 +165,12 @@ func TestUpgradeHandlerNewVersion(t *testing.T) {
 	ack := noopacker.New()
 	err1 := u.Handle(ctx, &a1, ack)
 	require.NoError(t, err1)
-	time.Sleep(1 * time.Second)
 	err2 := u.Handle(ctx, &a2, ack)
 	require.NoError(t, err2)
 	msg1 := <-msgChan
 	require.Equal(t, "canceled 8.2.0", msg1)
+	// indicate that upgrade is completed
+	close(completedChan)
 	msg2 := <-msgChan
 	require.Equal(t, "completed 8.5.0", msg2)
 }
diff --git a/testing/integration/agent_long_running_leak_test.go b/testing/integration/agent_long_running_leak_test.go
index f3862665771..6ccc567f55a 100644
--- a/testing/integration/agent_long_running_leak_test.go
+++ b/testing/integration/agent_long_running_leak_test.go
@@ -151,7 +151,7 @@ func (runner *ExtendedRunner) TestHandleLeak() {
 	timer := time.NewTimer(testDuration)
 	defer timer.Stop()
 
-	ticker := time.NewTicker(time.Second * 10)
+	ticker := time.NewTicker(time.Second * 60)
 	defer ticker.Stop()
 
 	done := false