WIP - Add Upgrader unit tests

Author: pchila

internal/pkg/agent/application/upgrade/upgrade.go

@@ -71,6 +71,14 @@ func init() {
 	}
 }
 
+// TODO substitute all the references to watcher with calls to the interface
+type WatcherHelper interface {
+	InvokeWatcher(log *logger.Logger, agentExecutable string) (*exec.Cmd, error)
+	SelectWatcherExecutable(topDir string, agentInstalls ...agentInstall) string
+	WaitForWatcher(ctx context.Context, log *logger.Logger, markerFilePath string, waitTime time.Duration) error
+	TakeOverWatcher(ctx context.Context, topDir string) (*filelock.AppLocker, error)
+}
+
 // Upgrader performs an upgrade
 type Upgrader struct {
 	log             *logger.Logger
@@ -204,7 +212,7 @@ func checkUpgrade(log *logger.Logger, currentVersion, newVersion agentVersion, m
 func (u *Upgrader) Upgrade(ctx context.Context, version string, rollback bool, sourceURI string, action *fleetapi.ActionUpgrade, det *details.Details, skipVerifyOverride bool, skipDefaultPgp bool, pgpBytes ...string) (_ reexec.ShutdownCallbackFn, err error) {
 
 	if rollback {
-		return u.forceRollbackToPreviousVersion(ctx, version, action, det)
+		return u.forceRollbackToPreviousVersion(ctx, paths.Top(), version, action, det)
 	}
 
 	u.log.Infow("Upgrading agent", "version", version, "source_uri", sourceURI)
@@ -388,19 +396,19 @@ func (u *Upgrader) Upgrade(ctx context.Context, version string, rollback bool, s
 	return cb, nil
 }
 
-func (u *Upgrader) forceRollbackToPreviousVersion(ctx context.Context, version string, action *fleetapi.ActionUpgrade, d *details.Details) (reexec.ShutdownCallbackFn, error) {
+func (u *Upgrader) forceRollbackToPreviousVersion(ctx context.Context, topDir string, version string, action *fleetapi.ActionUpgrade, d *details.Details) (reexec.ShutdownCallbackFn, error) {
 	// Formal checks for verifying we can rollback properly:
 	// 1. d.Metadata.RollbacksAvailable should contain the desired version with a valid TTL (it may need to be written by main agent process before starting watcher)
 	// 2. there has been at least the first restart with the new agent (i.e. we are not still downloading/extracting/rotating)
 	// 3. upgrade marker exists
 	// these should be revalidated after taking over watcher
-	err := u.PersistManualRollback(ctx)
+	err := u.PersistManualRollback(ctx, topDir)
 	if err != nil {
 		return nil, err
 	}
 
 	// Invoke watcher again
-	_, err = InvokeWatcher(u.log, paths.BinaryPath(paths.VersionedHome(paths.Top()), agentName))
+	_, err = InvokeWatcher(u.log, paths.BinaryPath(paths.VersionedHome(topDir), agentName))
 	if err != nil {
 		return nil, fmt.Errorf("invoking watcher: %w", err)
 	}
@@ -409,8 +417,8 @@ func (u *Upgrader) forceRollbackToPreviousVersion(ctx context.Context, version s
 
 }
 
-func (u *Upgrader) PersistManualRollback(ctx context.Context) error {
-	watcherApplock, err := u.takeOverWatcher(ctx)
+func (u *Upgrader) PersistManualRollback(ctx context.Context, topDir string) error {
+	watcherApplock, err := u.takeOverWatcher(ctx, topDir)
 	if err != nil {
 		return fmt.Errorf("taking over watcher processes: %w", err)
 	}
@@ -422,7 +430,7 @@ func (u *Upgrader) PersistManualRollback(ctx context.Context) error {
 	}(watcherApplock)
 
 	// read the upgrade marker
-	updateMarker, err := LoadMarker(paths.Data())
+	updateMarker, err := LoadMarker(paths.DataFrom(topDir))
 	if err != nil {
 		return fmt.Errorf("loading marker: %w", err)
 	}
@@ -435,7 +443,7 @@ func (u *Upgrader) PersistManualRollback(ctx context.Context) error {
 	return nil
 }
 
-func (u *Upgrader) takeOverWatcher(ctx context.Context) (*filelock.AppLocker, error) {
+func (u *Upgrader) takeOverWatcher(ctx context.Context, topDir string) (*filelock.AppLocker, error) {
 
 	takeoverCtx, takeoverCancel := context.WithTimeout(ctx, 30*time.Second)
 	defer takeoverCancel()
@@ -479,7 +487,7 @@ func (u *Upgrader) takeOverWatcher(ctx context.Context) (*filelock.AppLocker, er
 		case <-takeoverCtx.Done():
 			return nil, fmt.Errorf("timed out taking over watcher applocker")
 		case <-takeOverTicker.C:
-			locker := filelock.NewAppLocker(paths.Top(), "watcher.lock")
+			locker := filelock.NewAppLocker(topDir, "watcher.lock")
 			err := locker.TryLock()
 			if err != nil {
 				u.log.Errorf("error locking watcher applocker: %s", err)

internal/pkg/agent/application/upgrade/upgrade_test.go

@@ -26,6 +26,7 @@ import (
 	"github.com/elastic/elastic-agent/internal/pkg/agent/application/paths"
 	"github.com/elastic/elastic-agent/internal/pkg/agent/application/upgrade/artifact"
 	"github.com/elastic/elastic-agent/internal/pkg/agent/application/upgrade/details"
+	"github.com/elastic/elastic-agent/internal/pkg/agent/configuration"
 	"github.com/elastic/elastic-agent/internal/pkg/agent/errors"
 	"github.com/elastic/elastic-agent/internal/pkg/config"
 	"github.com/elastic/elastic-agent/internal/pkg/fleetapi"
@@ -36,6 +37,7 @@ import (
 	"github.com/elastic/elastic-agent/pkg/core/logger"
 	"github.com/elastic/elastic-agent/pkg/core/logger/loggertest"
 	agtversion "github.com/elastic/elastic-agent/pkg/version"
+	infomocks "github.com/elastic/elastic-agent/testing/mocks/internal_/pkg/agent/application/info"
 	ackermocks "github.com/elastic/elastic-agent/testing/mocks/internal_/pkg/fleetapi/acker"
 	clientmocks "github.com/elastic/elastic-agent/testing/mocks/pkg/control/v2/client"
 )
@@ -316,7 +318,7 @@ func TestUpgraderAckAction(t *testing.T) {
 		errAck := errors.New("ack error")
 		mockAcker.EXPECT().Ack(mock.Anything, action).Return(errAck)
 		// no expectation on Commit() since it shouldn't be called after an error during Ack()
-		
+
 		require.ErrorIs(t, u.AckAction(t.Context(), mockAcker, action), errAck)
 	})
 }
@@ -1270,3 +1272,12 @@ func TestIsSameReleaseVersion(t *testing.T) {
 		})
 	}
 }
+
+func TestManualRollback(t *testing.T) {
+	log, _ := loggertest.New(t.Name())
+	mockAgentInfo := infomocks.NewAgent(t)
+	upgrader, err := NewUpgrader(log, &artifact.Config{}, &configuration.UpgradeConfig{}, mockAgentInfo)
+	require.NoError(t, err, "error creating upgrader")
+
+	upgrader.Upgrade(t.Context(), "", true, "", nil, nil, true, false)
+}