test watcher command

Add some unit tests for watchCmd function, covering the most common
cases. This has required some refactoring to properly assert side
effects and interactions.

Author: pchila

.mockery.yaml

@@ -29,3 +29,11 @@ packages:
   github.com/elastic/elastic-agent/internal/pkg/agent/application/info:
     interfaces:
       Agent:
+  github.com/elastic/elastic-agent/internal/pkg/agent/cmd:
+    interfaces:
+      agentWatcher:
+        config:
+          mockname: "AgentWatcher"
+      installationModifier:
+        config:
+          mockname: "InstallationModifier"

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

@@ -313,8 +313,8 @@ func loadMarker(markerFile string) (*UpdateMarker, error) {
 // SaveMarker serializes and persists the given upgrade marker to disk.
 // For critical upgrade transitions, pass shouldFsync as true so the marker
 // file is immediately flushed to persistent storage.
-func SaveMarker(marker *UpdateMarker, shouldFsync bool) error {
-	return saveMarkerToPath(marker, markerFilePath(paths.Data()), shouldFsync)
+func SaveMarker(dataDirPath string, marker *UpdateMarker, shouldFsync bool) error {
+	return saveMarkerToPath(marker, markerFilePath(dataDirPath), shouldFsync)
 }
 
 func saveMarkerToPath(marker *UpdateMarker, markerFile string, shouldFsync bool) error {

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

@@ -449,7 +449,7 @@ func (u *Upgrader) Ack(ctx context.Context, acker acker.Acker) error {
 
 	marker.Acked = true
 
-	return SaveMarker(marker, false)
+	return SaveMarker(paths.Data(), marker, false)
 }
 
 func (u *Upgrader) AckAction(ctx context.Context, acker acker.Acker, action fleetapi.Action) error {

internal/pkg/agent/cmd/watch.go

@@ -53,7 +53,7 @@ func newWatchCommandWithArgs(_ []string, streams *cli.IOStreams) *cobra.Command
 			// Make sure to flush any buffered logs before we're done.
 			defer log.Sync() //nolint:errcheck // flushing buffered logs is best effort.
 
-			if err := watchCmd(log, cfg); err != nil {
+			if err := watchCmd(log, paths.Top(), cfg.Settings.Upgrade.Watcher, new(agtWatcherImpl), new(agtInstallationModifierImpl)); err != nil {
 				log.Errorw("Watch command failed", "error.message", err)
 				fmt.Fprintf(streams.Err, "Watch command failed: %v\n%s\n", err, troubleshootMessage())
 				os.Exit(4)
@@ -64,21 +64,30 @@ func newWatchCommandWithArgs(_ []string, streams *cli.IOStreams) *cobra.Command
 	return cmd
 }
 
-func watchCmd(log *logp.Logger, cfg *configuration.Configuration) error {
+type agentWatcher interface {
+	Watch(ctx context.Context, tilGrace, errorCheckInterval time.Duration, log *logp.Logger) error
+}
+
+type installationModifier interface {
+	Cleanup(log *logger.Logger, topDirPath, currentVersionedHome, currentHash string, removeMarker, keepLogs bool) error
+	Rollback(ctx context.Context, log *logger.Logger, c client.Client, topDirPath, prevVersionedHome, prevHash string) error
+}
+
+func watchCmd(log *logp.Logger, topDir string, cfg *configuration.UpgradeWatcherConfig, watcher agentWatcher, installModifier installationModifier) error {
 	log.Infow("Upgrade Watcher started", "process.pid", os.Getpid(), "agent.version", version.GetAgentPackageVersion())
-	marker, err := upgrade.LoadMarker(paths.Data())
+	dataDir := paths.DataFrom(topDir)
+	marker, err := upgrade.LoadMarker(dataDir)
 	if err != nil {
 		log.Error("failed to load marker", err)
 		return err
 	}
 	if marker == nil {
 		// no marker found we're not in upgrade process
-		log.Infof("update marker not present at '%s'", paths.Data())
+		log.Infof("update marker not present at '%s'", dataDir)
 		return nil
 	}
 
-	log.Infof("Loaded update marker %+v", marker)
-
+	log.With("marker", marker, "details", marker.Details).Info("Loaded update marker")
 	locker := filelock.NewAppLocker(paths.Top(), watcherLockFile)
 	if err := locker.TryLock(); err != nil {
 		if errors.Is(err, filelock.ErrAppAlreadyRunning) {
@@ -93,14 +102,18 @@ func watchCmd(log *logp.Logger, cfg *configuration.Configuration) error {
 		_ = locker.Unlock()
 	}()
 
-	isWithinGrace, tilGrace := gracePeriod(marker, cfg.Settings.Upgrade.Watcher.GracePeriod)
-	if !isWithinGrace || isTerminalState(marker) {
-		log.Infof("not within grace [updatedOn %v] %v or agent have been rolled back [state: %s]", marker.UpdatedOn.String(), time.Since(marker.UpdatedOn).String(), marker.Details.State)
+	isWithinGrace, tilGrace := gracePeriod(marker, cfg.GracePeriod)
+	if isTerminalState(marker) || !isWithinGrace {
+		stateString := ""
+		if marker.Details != nil {
+			stateString = string(marker.Details.State)
+		}
+		log.Infof("not within grace [updatedOn %v] %v or agent have been rolled back [state: %s]", marker.UpdatedOn.String(), time.Since(marker.UpdatedOn).String(), stateString)
 		// if it is started outside of upgrade loop
 		// if we're not within grace and marker is still there it might mean
 		// that cleanup was not performed ok, cleanup everything except current version
 		// hash is the same as hash of agent which initiated watcher.
-		if err := upgrade.Cleanup(log, paths.Top(), paths.VersionedHome(paths.Top()), release.ShortCommit(), true, false); err != nil {
+		if err := installModifier.Cleanup(log, paths.Top(), paths.VersionedHome(topDir), release.ShortCommit(), true, false); err != nil {
 			log.Error("clean up of prior watcher run failed", err)
 		}
 		// exit nicely
@@ -109,15 +122,18 @@ func watchCmd(log *logp.Logger, cfg *configuration.Configuration) error {
 
 	// About to start watching the upgrade. Initialize upgrade details and save them in the
 	// upgrade marker.
-	upgradeDetails := initUpgradeDetails(marker, upgrade.SaveMarker, log)
+	saveMarkerFunc := func(marker *upgrade.UpdateMarker, b bool) error {
+		return upgrade.SaveMarker(dataDir, marker, b)
+	}
+	upgradeDetails := initUpgradeDetails(marker, saveMarkerFunc, log)
 
-	errorCheckInterval := cfg.Settings.Upgrade.Watcher.ErrorCheck.Interval
+	errorCheckInterval := cfg.ErrorCheck.Interval
 	ctx := context.Background()
-	if err := watch(ctx, tilGrace, errorCheckInterval, log); err != nil {
+	if err := watcher.Watch(ctx, tilGrace, errorCheckInterval, log); err != nil {
 		log.Error("Error detected, proceeding to rollback: %v", err)
 
 		upgradeDetails.SetStateWithReason(details.StateRollback, "automatic rollback")
-		err = upgrade.Rollback(ctx, log, client.New(), paths.Top(), marker.PrevVersionedHome, marker.PrevHash)
+		err = installModifier.Rollback(ctx, log, client.New(), paths.Top(), marker.PrevVersionedHome, marker.PrevHash)
 		if err != nil {
 			log.Error("rollback failed", err)
 			upgradeDetails.Fail(err)
@@ -135,7 +151,7 @@ func watchCmd(log *logp.Logger, cfg *configuration.Configuration) error {
 	// Why is this being skipped on Windows? The comment above is not clear.
 	// issue: https://github.com/elastic/elastic-agent/issues/3027
 	removeMarker := !isWindows()
-	err = upgrade.Cleanup(log, paths.Top(), marker.VersionedHome, marker.Hash, removeMarker, false)
+	err = installModifier.Cleanup(log, topDir, marker.VersionedHome, marker.Hash, removeMarker, false)
 	if err != nil {
 		log.Error("cleanup after successful watch failed", err)
 	}

internal/pkg/agent/cmd/watch_impl.go

@@ -0,0 +1,31 @@
+// Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+// or more contributor license agreements. Licensed under the Elastic License 2.0;
+// you may not use this file except in compliance with the Elastic License 2.0.
+
+package cmd
+
+import (
+	"context"
+	"time"
+
+	"github.com/elastic/elastic-agent-libs/logp"
+	"github.com/elastic/elastic-agent/internal/pkg/agent/application/upgrade"
+	"github.com/elastic/elastic-agent/pkg/control/v2/client"
+	"github.com/elastic/elastic-agent/pkg/core/logger"
+)
+
+type agtWatcherImpl struct{}
+
+func (a agtWatcherImpl) Watch(ctx context.Context, tilGrace, errorCheckInterval time.Duration, log *logp.Logger) error {
+	return watch(ctx, tilGrace, errorCheckInterval, log)
+}
+
+type agtInstallationModifierImpl struct{}
+
+func (a agtInstallationModifierImpl) Cleanup(log *logger.Logger, topDirPath, currentVersionedHome, currentHash string, removeMarker, keepLogs bool) error {
+	return upgrade.Cleanup(log, topDirPath, currentVersionedHome, currentHash, removeMarker, keepLogs)
+}
+
+func (a agtInstallationModifierImpl) Rollback(ctx context.Context, log *logger.Logger, c client.Client, topDirPath, prevVersionedHome, prevHash string) error {
+	return upgrade.Rollback(ctx, log, c, topDirPath, prevVersionedHome, prevHash)
+}

internal/pkg/agent/cmd/watch_test.go

@@ -5,18 +5,28 @@
 package cmd
 
 import (
+	"fmt"
+	"os"
+	"runtime"
 	"testing"
+	"time"
 
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/mock"
 	"go.uber.org/zap/zapcore"
 
+	"github.com/stretchr/testify/require"
+
+	"github.com/elastic/elastic-agent/internal/pkg/agent/application/paths"
 	"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/fleetapi"
+	"github.com/elastic/elastic-agent/internal/pkg/release"
 	"github.com/elastic/elastic-agent/pkg/core/logger/loggertest"
 
-	"github.com/stretchr/testify/require"
-
 	"github.com/elastic/elastic-agent/internal/pkg/agent/application/upgrade"
+	cmdmocks "github.com/elastic/elastic-agent/testing/mocks/internal_/pkg/agent/cmd"
 )
 
 func TestInitUpgradeDetails(t *testing.T) {
@@ -74,3 +84,202 @@ func TestInitUpgradeDetails(t *testing.T) {
 	require.Equal(t, zapcore.ErrorLevel, logs[0].Level)
 	require.Equal(t, `unable to save upgrade marker after clearing upgrade details: some error`, logs[0].Message)
 }
+
+func Test_watchCmd(t *testing.T) {
+	type args struct {
+		cfg *configuration.UpgradeWatcherConfig
+	}
+	tests := []struct {
+		name               string
+		setupUpgradeMarker func(t *testing.T, tmpDir string, watcher *cmdmocks.AgentWatcher, installModifier *cmdmocks.InstallationModifier)
+		args               args
+		wantErr            assert.ErrorAssertionFunc
+	}{
+		{
+			name: "no upgrade marker, no party",
+			setupUpgradeMarker: func(t *testing.T, topDir string, watcher *cmdmocks.AgentWatcher, installModifier *cmdmocks.InstallationModifier) {
+				dataDirPath := paths.DataFrom(topDir)
+				err := os.MkdirAll(dataDirPath, 0755)
+				require.NoError(t, err)
+			},
+			args: args{
+				cfg: configuration.DefaultUpgradeConfig().Watcher,
+			},
+			wantErr: assert.NoError,
+		},
+		{
+			name: "happy path: no error watching, cleanup prev install",
+			setupUpgradeMarker: func(t *testing.T, topDir string, watcher *cmdmocks.AgentWatcher, installModifier *cmdmocks.InstallationModifier) {
+				dataDirPath := paths.DataFrom(topDir)
+				err := os.MkdirAll(dataDirPath, 0755)
+				require.NoError(t, err)
+				err = upgrade.SaveMarker(
+					dataDirPath,
+					&upgrade.UpdateMarker{
+						Version:           "4.5.6",
+						Hash:              "newver",
+						VersionedHome:     "elastic-agent-4.5.6-newver",
+						UpdatedOn:         time.Now(),
+						PrevVersion:       "1.2.3",
+						PrevHash:          "prvver",
+						PrevVersionedHome: "elastic-agent-prvver",
+						Acked:             false,
+						Action:            nil,
+						Details:           nil, //details.NewDetails("4.5.6", details.StateReplacing, ""),
+						DesiredOutcome:    upgrade.OUTCOME_UPGRADE,
+					},
+					true,
+				)
+				require.NoError(t, err)
+
+				watcher.EXPECT().
+					Watch(mock.Anything, mock.Anything, mock.Anything, mock.Anything).
+					Return(nil)
+
+				expectedRemoveMarkerFlag := true
+				if runtime.GOOS == "windows" {
+					// on windows the marker is not removed immediately to allow for cleanup on restart
+					expectedRemoveMarkerFlag = false
+				}
+				installModifier.EXPECT().
+					Cleanup(mock.Anything, topDir, "elastic-agent-4.5.6-newver", "newver", expectedRemoveMarkerFlag, false).
+					Return(nil)
+			},
+			args: args{
+				cfg: configuration.DefaultUpgradeConfig().Watcher,
+			},
+			wantErr: assert.NoError,
+		},
+		{
+			name: "unhappy path: error watching, rollback to previous install",
+			setupUpgradeMarker: func(t *testing.T, topDir string, watcher *cmdmocks.AgentWatcher, installModifier *cmdmocks.InstallationModifier) {
+				dataDirPath := paths.DataFrom(topDir)
+				err := os.MkdirAll(dataDirPath, 0755)
+				require.NoError(t, err)
+				err = upgrade.SaveMarker(
+					dataDirPath,
+					&upgrade.UpdateMarker{
+						Version:           "4.5.6",
+						Hash:              "newver",
+						VersionedHome:     "elastic-agent-4.5.6-newver",
+						UpdatedOn:         time.Now(),
+						PrevVersion:       "1.2.3",
+						PrevHash:          "prvver",
+						PrevVersionedHome: "elastic-agent-prvver",
+						Acked:             false,
+						Action:            nil,
+						Details:           nil, //details.NewDetails("4.5.6", details.StateReplacing, ""),
+						DesiredOutcome:    upgrade.OUTCOME_UPGRADE,
+					},
+					true,
+				)
+				require.NoError(t, err)
+
+				watcher.EXPECT().
+					Watch(mock.Anything, mock.Anything, mock.Anything, mock.Anything).
+					Return(errors.New("some watch error due to agent misbehaving"))
+				installModifier.EXPECT().
+					Rollback(mock.Anything, mock.Anything, mock.Anything, paths.Top(), "elastic-agent-prvver", "prvver").
+					Return(nil)
+			},
+			args: args{
+				cfg: configuration.DefaultUpgradeConfig().Watcher,
+			},
+			wantErr: assert.NoError,
+		},
+		{
+			name: "upgrade rolled back: no watching, cleanup must be called",
+			setupUpgradeMarker: func(t *testing.T, topDir string, watcher *cmdmocks.AgentWatcher, installModifier *cmdmocks.InstallationModifier) {
+				dataDirPath := paths.DataFrom(topDir)
+				err := os.MkdirAll(dataDirPath, 0755)
+				require.NoError(t, err)
+				err = upgrade.SaveMarker(
+					dataDirPath,
+					&upgrade.UpdateMarker{
+						Version:           "4.5.6",
+						Hash:              "newver",
+						VersionedHome:     "elastic-agent-4.5.6-newver",
+						UpdatedOn:         time.Now(),
+						PrevVersion:       "1.2.3",
+						PrevHash:          "prvver",
+						PrevVersionedHome: "elastic-agent-prvver",
+						Acked:             false,
+						Action:            nil,
+						Details: &details.Details{
+							TargetVersion: "4.5.6",
+							State:         details.StateRollback,
+							Metadata: details.Metadata{
+								Reason: "automatic rollback",
+							},
+						},
+						DesiredOutcome: upgrade.OUTCOME_UPGRADE,
+					},
+					true,
+				)
+				require.NoError(t, err)
+				// topdir, prevVersionedHome and prevHash are not taken from the upgrade marker, otherwise they would be
+				// <topDir, "topDir/data/elastic-agent-prvver", "prvver">
+				installModifier.EXPECT().
+					Cleanup(mock.Anything, paths.Top(), paths.VersionedHome(topDir), release.ShortCommit(), true, false).
+					Return(nil)
+			},
+			args: args{
+				cfg: configuration.DefaultUpgradeConfig().Watcher,
+			},
+			wantErr: assert.NoError,
+		},
+		{
+			name: "after grace period: no watching, cleanup must be called",
+			setupUpgradeMarker: func(t *testing.T, topDir string, watcher *cmdmocks.AgentWatcher, installModifier *cmdmocks.InstallationModifier) {
+				dataDirPath := paths.DataFrom(topDir)
+				err := os.MkdirAll(dataDirPath, 0755)
+				require.NoError(t, err)
+				updatedOn := time.Now().Add(-5 * time.Minute)
+				err = upgrade.SaveMarker(
+					dataDirPath,
+					&upgrade.UpdateMarker{
+						Version:           "4.5.6",
+						Hash:              "newver",
+						VersionedHome:     "elastic-agent-4.5.6-newver",
+						UpdatedOn:         updatedOn,
+						PrevVersion:       "1.2.3",
+						PrevHash:          "prvver",
+						PrevVersionedHome: "elastic-agent-prvver",
+						Acked:             false,
+						Action:            nil,
+						Details:           nil,
+						DesiredOutcome:    upgrade.OUTCOME_UPGRADE,
+					},
+					true,
+				)
+				require.NoError(t, err)
+
+				// topdir, prevVersionedHome and prevHash are not taken from the upgrade marker, otherwise they would be
+				// <topDir, "topDir/data/elastic-agent-prvver", "prvver">
+				installModifier.EXPECT().
+					Cleanup(mock.Anything, paths.Top(), paths.VersionedHome(topDir), release.ShortCommit(), true, false).
+					Return(nil)
+			},
+			args: args{
+				cfg: &configuration.UpgradeWatcherConfig{
+					GracePeriod: 2 * time.Minute,
+					ErrorCheck: configuration.UpgradeWatcherCheckConfig{
+						Interval: time.Second,
+					},
+				},
+			},
+			wantErr: assert.NoError,
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			log, obs := loggertest.New(t.Name())
+			tmpDir := t.TempDir()
+			mockWatcher := cmdmocks.NewAgentWatcher(t)
+			mockInstallModifier := cmdmocks.NewInstallationModifier(t)
+			tt.setupUpgradeMarker(t, tmpDir, mockWatcher, mockInstallModifier)
+			tt.wantErr(t, watchCmd(log, tmpDir, tt.args.cfg, mockWatcher, mockInstallModifier), fmt.Sprintf("watchCmd(%v, ...)", tt.args.cfg))
+			t.Logf("watchCmd logs:\n%v", obs.All())
+		})
+	}
+}