Fix flakiness in TestUpgradeHandler*

Fix the flakiness from TestUpgradeHandler* tests by re-working the
mockUpgradeManager, now it accepts a function for its Upgrade method
and their implementation is goroutine safe

Author: belimawr

internal/pkg/agent/application/actions/handlers/handler_action_upgrade_test.go

@@ -7,7 +7,9 @@ package handlers
 import (
 	"context"
 	"errors"
+	"sync/atomic"
 	"testing"
+	"time"
 
 	"github.com/stretchr/testify/require"
 
@@ -27,9 +29,15 @@ import (
 )
 
 type mockUpgradeManager struct {
-	msgChan     chan string
-	successChan chan string
-	errChan     chan error
+	UpgradeFn func(
+		ctx context.Context,
+		version string,
+		sourceURI string,
+		action *fleetapi.ActionUpgrade,
+		details *details.Details,
+		skipVerifyOverride bool,
+		skipDefaultPgp bool,
+		pgpBytes ...string) (reexec.ShutdownCallbackFn, error)
 }
 
 func (u *mockUpgradeManager) Upgradeable() bool {
@@ -40,17 +48,25 @@ func (u *mockUpgradeManager) Reload(rawConfig *config.Config) error {
 	return nil
 }
 
-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 msg := <-u.successChan:
-		u.msgChan <- msg
-		return nil, nil
-	case err := <-u.errChan:
-		u.msgChan <- err.Error()
-		return nil, ctx.Err()
-	case <-ctx.Done():
-		return nil, ctx.Err()
-	}
+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, error) {
+
+	return u.UpgradeFn(
+		ctx,
+		version,
+		sourceURI,
+		action,
+		details,
+		skipVerifyOverride,
+		skipDefaultPgp,
+		pgpBytes...)
 }
 
 func (u *mockUpgradeManager) Ack(ctx context.Context, acker acker.Acker) error {
@@ -70,8 +86,7 @@ func TestUpgradeHandler(t *testing.T) {
 	log, _ := logger.New("", false)
 
 	agentInfo := &info.AgentInfo{}
-	msgChan := make(chan string)
-	completedChan := make(chan string)
+	upgradeCalledChan := make(chan struct{})
 
 	// Create and start the coordinator
 	c := coordinator.New(
@@ -81,7 +96,21 @@ func TestUpgradeHandler(t *testing.T) {
 		agentInfo,
 		component.RuntimeSpecs{},
 		nil,
-		&mockUpgradeManager{msgChan: msgChan, successChan: completedChan},
+		&mockUpgradeManager{
+			UpgradeFn: func(
+				ctx context.Context,
+				version string,
+				sourceURI string,
+				action *fleetapi.ActionUpgrade,
+				details *details.Details,
+				skipVerifyOverride bool,
+				skipDefaultPgp bool,
+				pgpBytes ...string) (reexec.ShutdownCallbackFn, error) {
+
+				upgradeCalledChan <- struct{}{}
+				return nil, nil
+			},
+		},
 		nil, nil, nil, nil, nil, false)
 	//nolint:errcheck // We don't need the termination state of the Coordinator
 	go c.Run(ctx)
@@ -91,11 +120,14 @@ 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)
+
+	// Make sure this test does not dead lock or wait for too long
+	select {
+	case <-time.Tick(50 * time.Millisecond):
+		t.Fatal("mockUpgradeManager.Upgrade was not called")
+	case <-upgradeCalledChan:
+	}
 }
 
 func TestUpgradeHandlerSameVersion(t *testing.T) {
@@ -109,19 +141,37 @@ func TestUpgradeHandlerSameVersion(t *testing.T) {
 	log, _ := logger.New("", false)
 
 	agentInfo := &info.AgentInfo{}
-	msgChan := make(chan string)
-	successChan := make(chan string)
-	errChan := make(chan error)
+	upgradeCalledChan := make(chan struct{})
 
 	// Create and start the Coordinator
+	upgradeCalled := atomic.Bool{}
 	c := coordinator.New(
 		log,
 		configuration.DefaultConfiguration(),
 		logger.DefaultLogLevel,
 		agentInfo,
 		component.RuntimeSpecs{},
 		nil,
-		&mockUpgradeManager{msgChan: msgChan, successChan: successChan, errChan: errChan},
+		&mockUpgradeManager{
+			UpgradeFn: func(
+				ctx context.Context,
+				version string,
+				sourceURI string,
+				action *fleetapi.ActionUpgrade,
+				details *details.Details,
+				skipVerifyOverride bool,
+				skipDefaultPgp bool,
+				pgpBytes ...string) (reexec.ShutdownCallbackFn, error) {
+
+				if upgradeCalled.CompareAndSwap(false, true) {
+					upgradeCalledChan <- struct{}{}
+					return nil, nil
+				}
+				err := errors.New("mockUpgradeManager.Upgrade called more than once")
+				t.Error(err.Error())
+				return nil, err
+			},
+		},
 		nil, nil, nil, nil, nil, false)
 	//nolint:errcheck // We don't need the termination state of the Coordinator
 	go c.Run(ctx)
@@ -135,11 +185,12 @@ func TestUpgradeHandlerSameVersion(t *testing.T) {
 	require.NoError(t, err1)
 	require.NoError(t, err2)
 
-	successChan <- "completed 8.3.0"
-	require.Equal(t, "completed 8.3.0", <-msgChan)
-	errChan <- errors.New("duplicated update, not finishing it?")
-	require.Equal(t, "duplicated update, not finishing it?", <-msgChan)
-
+	// Make sure this test does not dead lock or wait for too long
+	select {
+	case <-time.Tick(50 * time.Millisecond):
+		t.Fatal("mockUpgradeManager.Upgrade was not called")
+	case <-upgradeCalledChan:
+	}
 }
 
 func TestUpgradeHandlerNewVersion(t *testing.T) {
@@ -149,11 +200,9 @@ func TestUpgradeHandlerNewVersion(t *testing.T) {
 	defer cancel()
 
 	log, _ := logger.New("", false)
+	upgradeCalledChan := make(chan string)
 
 	agentInfo := &info.AgentInfo{}
-	msgChan := make(chan string)
-	completedChan := make(chan string)
-	errorChan := make(chan error)
 
 	// Create and start the Coordinator
 	c := coordinator.New(
@@ -163,7 +212,27 @@ func TestUpgradeHandlerNewVersion(t *testing.T) {
 		agentInfo,
 		component.RuntimeSpecs{},
 		nil,
-		&mockUpgradeManager{msgChan: msgChan, successChan: completedChan, errChan: errorChan},
+		&mockUpgradeManager{
+			UpgradeFn: func(
+				ctx context.Context,
+				version string,
+				sourceURI string,
+				action *fleetapi.ActionUpgrade,
+				details *details.Details,
+				skipVerifyOverride bool,
+				skipDefaultPgp bool,
+				pgpBytes ...string) (reexec.ShutdownCallbackFn, error) {
+
+				defer func() {
+					upgradeCalledChan <- version
+				}()
+				if version == "8.2.0" {
+					return nil, errors.New("upgrade to 8.2.0 will always fail")
+				}
+
+				return nil, nil
+			},
+		},
 		nil, nil, nil, nil, nil, false)
 	//nolint:errcheck // We don't need the termination state of the Coordinator
 	go c.Run(ctx)
@@ -175,18 +244,24 @@ func TestUpgradeHandlerNewVersion(t *testing.T) {
 		Version: "8.5.0", SourceURI: "http://localhost"}}
 	ack := noopacker.New()
 
+	checkMsg := func(c <-chan string, expected, errMsg string) {
+		t.Helper()
+		// Make sure this test does not dead lock or wait for too long
+		// For some reason < 1s sometimes makes the test fail.
+		select {
+		case <-time.Tick(1300 * time.Millisecond):
+			t.Fatal("timed out waiting for Upgrade to return")
+		case msg := <-c:
+			require.Equal(t, expected, msg, errMsg)
+		}
+	}
+
 	// Send both upgrade actions, a1 will error before a2 succeeds
 	err1 := u.Handle(ctx, &a1, ack)
 	require.NoError(t, err1)
+	checkMsg(upgradeCalledChan, "8.2.0", "first call must be with version 8.2.0")
+
 	err2 := u.Handle(ctx, &a2, ack)
 	require.NoError(t, err2)
-
-	// Send an error so the first action is "cancelled"
-	errorChan <- errors.New("cancelled 8.2.0")
-	// Wait for the mockUpgradeHandler to receive and "process" the error
-	require.Equal(t, "cancelled 8.2.0", <-msgChan, "mockUpgradeHandler.Upgrade did not receive the expected error")
-
-	// Send a success so the second action succeeds
-	completedChan <- "completed 8.5.0"
-	require.Equal(t, "completed 8.5.0", <-msgChan, "mockUpgradeHandler.Upgrade did not receive the success signal")
+	checkMsg(upgradeCalledChan, "8.5.0", "second call to Upgrade must be with version 8.5.0")
 }