fixup! Implement manual rollback from list of agent installs

Author: pchila

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

@@ -5,7 +5,7 @@
 package upgrade
 
 import (
-	"io/fs"
+	"fmt"
 	"os"
 	"os/exec"
 	"path/filepath"
@@ -15,14 +15,19 @@ import (
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/mock"
 	"github.com/stretchr/testify/require"
+	"gopkg.in/yaml.v3"
 
 	"github.com/elastic/elastic-agent/internal/pkg/agent/application/filelock"
 	"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/release"
+	v1 "github.com/elastic/elastic-agent/pkg/api/v1"
 	"github.com/elastic/elastic-agent/pkg/core/logger/loggertest"
 	"github.com/elastic/elastic-agent/pkg/version"
 	"github.com/elastic/elastic-agent/testing/mocks/internal_/pkg/agent/application/info"
+	agtversion "github.com/elastic/elastic-agent/version"
 )
 
 func TestManualRollback(t *testing.T) {
@@ -84,14 +89,28 @@ func TestManualRollback(t *testing.T) {
 		versionedHome: "data/elastic-agent-4.5.6-newver",
 	}
 
+	agentInstallCurrent := agentInstall{
+		parsedVersion: agtversion.GetParsedAgentPackageVersion(),
+		version:       release.VersionWithSnapshot(),
+		hash:          release.Commit(),
+		// Versioned home should contain the version but since the path does not really exist we fallback to the legacy format with just the hash
+		// versionedHome: filepath.Join("data", fmt.Sprintf("elastic-agent-%s-%s", release.VersionWithSnapshot(), release.ShortCommit())),
+		versionedHome: filepath.Join("data", fmt.Sprintf("elastic-agent-%s", release.ShortCommit())),
+	}
+
 	// this is the updated_on timestamp in the example
 	nowBeforeTTL, err := time.Parse(time.RFC3339, `2025-07-11T10:11:12Z`)
 	require.NoError(t, err, "error parsing nowBeforeTTL")
 
 	// the update marker yaml assume 7d TLL for rollbacks, let's make an extra day pass
 	nowAfterTTL := nowBeforeTTL.Add(8 * 24 * time.Hour)
 
-	type setupF func(t *testing.T, topDir string, agent *info.Agent, watcherHelper *MockWatcherHelper)
+	// save the current timestamp, useful for TTL-based testing
+	aMomentInTime := time.Now()
+	aMomentTomorrow := aMomentInTime.Add(24 * time.Hour)
+	aMomentAgo := aMomentInTime.Add(-1 * time.Second)
+
+	type setupF func(t *testing.T, topDir string, agent *info.Agent, watcherHelper *MockWatcherHelper, installSource *mockInstallDescriptorSource)
 	type postRollbackAssertionsF func(t *testing.T, topDir string)
 	type testcase struct {
 		name              string
@@ -107,7 +126,7 @@ func TestManualRollback(t *testing.T) {
 	testcases := []testcase{
 		{
 			name: "no rollback version - rollback fails",
-			setup: func(t *testing.T, topDir string, agent *info.Agent, watcherHelper *MockWatcherHelper) {
+			setup: func(t *testing.T, topDir string, agent *info.Agent, watcherHelper *MockWatcherHelper, installSource *mockInstallDescriptorSource) {
 				//do not setup anything here, let the rollback fail
 			},
 			artifactSettings: artifact.DefaultConfig(),
@@ -119,21 +138,34 @@ func TestManualRollback(t *testing.T) {
 			additionalAsserts: nil,
 		},
 		{
-			name: "no update marker - rollback fails",
-			setup: func(t *testing.T, topDir string, agent *info.Agent, watcherHelper *MockWatcherHelper) {
-				//do not setup anything here, let the rollback fail
+			name: "no update marker, no other installs - rollback fails",
+			setup: func(t *testing.T, topDir string, agent *info.Agent, watcherHelper *MockWatcherHelper, installSource *mockInstallDescriptorSource) {
+				installSource.EXPECT().GetInstallDesc().Return(
+					&v1.InstallDescriptor{
+						AgentInstalls: []v1.AgentInstallDesc{
+							{
+								Version:       "1.2.3",
+								Hash:          "oldver",
+								VersionedHome: "data/elastic-agent-1.2.3-oldver",
+								Flavor:        "strawberry",
+								Active:        true,
+							},
+						},
+					},
+					nil,
+				)
 			},
 			artifactSettings: artifact.DefaultConfig(),
 			upgradeSettings:  configuration.DefaultUpgradeConfig(),
 			version:          "1.2.3",
 			wantErr: func(t assert.TestingT, err error, i ...interface{}) bool {
-				return assert.ErrorIs(t, err, fs.ErrNotExist)
+				return assert.ErrorIs(t, err, ErrNoRollbacksAvailable)
 			},
 			additionalAsserts: nil,
 		},
 		{
 			name: "update marker is malformed - rollback fails",
-			setup: func(t *testing.T, topDir string, agent *info.Agent, watcherHelper *MockWatcherHelper) {
+			setup: func(t *testing.T, topDir string, agent *info.Agent, watcherHelper *MockWatcherHelper, installSource *mockInstallDescriptorSource) {
 				err := os.WriteFile(markerFilePath(paths.DataFrom(topDir)), []byte("this is not a proper YAML file"), 0600)
 				require.NoError(t, err, "error setting up update marker")
 				locker := filelock.NewAppLocker(topDir, "watcher.lock")
@@ -153,7 +185,7 @@ func TestManualRollback(t *testing.T) {
 		},
 		{
 			name: "update marker ok but rollback available is empty - error",
-			setup: func(t *testing.T, topDir string, agent *info.Agent, watcherHelper *MockWatcherHelper) {
+			setup: func(t *testing.T, topDir string, agent *info.Agent, watcherHelper *MockWatcherHelper, installSource *mockInstallDescriptorSource) {
 				err := os.WriteFile(markerFilePath(paths.DataFrom(topDir)), []byte(updatemarkerwatching456NoRollbackAvailable), 0600)
 				require.NoError(t, err, "error setting up update marker")
 				locker := filelock.NewAppLocker(topDir, "watcher.lock")
@@ -182,7 +214,7 @@ func TestManualRollback(t *testing.T) {
 		},
 		{
 			name: "update marker ok but version is not available for rollback - error",
-			setup: func(t *testing.T, topDir string, agent *info.Agent, watcherHelper *MockWatcherHelper) {
+			setup: func(t *testing.T, topDir string, agent *info.Agent, watcherHelper *MockWatcherHelper, installSource *mockInstallDescriptorSource) {
 				err := os.WriteFile(markerFilePath(paths.DataFrom(topDir)), []byte(updatemarkerwatching456), 0600)
 				require.NoError(t, err, "error setting up update marker")
 				locker := filelock.NewAppLocker(topDir, "watcher.lock")
@@ -211,7 +243,7 @@ func TestManualRollback(t *testing.T) {
 		},
 		{
 			name: "update marker ok but rollback is expired - error",
-			setup: func(t *testing.T, topDir string, agent *info.Agent, watcherHelper *MockWatcherHelper) {
+			setup: func(t *testing.T, topDir string, agent *info.Agent, watcherHelper *MockWatcherHelper, installSource *mockInstallDescriptorSource) {
 				err := os.WriteFile(markerFilePath(paths.DataFrom(topDir)), []byte(updatemarkerwatching456), 0600)
 				require.NoError(t, err, "error setting up update marker")
 				locker := filelock.NewAppLocker(topDir, "watcher.lock")
@@ -241,7 +273,7 @@ func TestManualRollback(t *testing.T) {
 		},
 		{
 			name: "update marker ok - takeover watcher, persist rollback and restart most recent watcher",
-			setup: func(t *testing.T, topDir string, agent *info.Agent, watcherHelper *MockWatcherHelper) {
+			setup: func(t *testing.T, topDir string, agent *info.Agent, watcherHelper *MockWatcherHelper, installSource *mockInstallDescriptorSource) {
 				err := os.WriteFile(markerFilePath(paths.DataFrom(topDir)), []byte(updatemarkerwatching456), 0600)
 				require.NoError(t, err, "error setting up update marker")
 				locker := filelock.NewAppLocker(topDir, "watcher.lock")
@@ -268,6 +300,120 @@ func TestManualRollback(t *testing.T) {
 				assert.NotEmpty(t, marker.RollbacksAvailable)
 			},
 		},
+		{
+			name: "no update marker, available install for rollback with valid TTL",
+			setup: func(t *testing.T, topDir string, agent *info.Agent, watcherHelper *MockWatcherHelper, installSource *mockInstallDescriptorSource) {
+				installSource.EXPECT().GetInstallDesc().Return(
+					&v1.InstallDescriptor{
+						AgentInstalls: []v1.AgentInstallDesc{
+							{
+								Version: release.VersionWithSnapshot(),
+								Hash:    release.Commit(),
+								// Versioned home should contain the version but since the path does not really exist we fallback to the legacy format with just the hash
+								//VersionedHome: filepath.Join("data", fmt.Sprintf("elastic-agent-%s-%s", release.VersionWithSnapshot(), release.ShortCommit())),
+								VersionedHome: filepath.Join("data", fmt.Sprintf("elastic-agent-%s", release.ShortCommit())),
+								Flavor:        "basic",
+								Active:        true,
+							},
+							{
+								// old install is still valid for the next 24  hours
+								OptionalTTLItem: v1.OptionalTTLItem{TTL: &aMomentTomorrow},
+								Version:         "1.2.3",
+								Hash:            "oldver",
+								VersionedHome:   "data/elastic-agent-1.2.3-oldver",
+								Flavor:          "basic",
+								Active:          false,
+							},
+						},
+					},
+					nil,
+				)
+				newerWatcherExecutable := filepath.Join(topDir, "data", fmt.Sprintf("elastic-agent-%s-%s", release.VersionWithSnapshot(), release.ShortCommit()), "elastic-agent")
+				watcherHelper.EXPECT().SelectWatcherExecutable(topDir, agentInstall123, agentInstallCurrent).Return(newerWatcherExecutable)
+				watcherHelper.EXPECT().InvokeWatcher(mock.Anything, newerWatcherExecutable, "watch", "--rollback", "data/elastic-agent-1.2.3-oldver").
+					Return(&exec.Cmd{Path: newerWatcherExecutable, Args: []string{"watch", "for rollbacksies"}, Process: &os.Process{Pid: 123}}, nil)
+			},
+			artifactSettings: artifact.DefaultConfig(),
+			upgradeSettings: &configuration.UpgradeConfig{
+				Rollback: &configuration.UpgradeRollbackConfig{
+					Window: 24 * time.Hour,
+				},
+			},
+			now:     aMomentInTime,
+			version: "1.2.3",
+			wantErr: assert.NoError,
+			additionalAsserts: func(t *testing.T, topDir string) {
+				actualMarkerFilePath := filepath.Join(topDir, "data", markerFilename)
+				require.FileExists(t, actualMarkerFilePath, "marker file must have been created")
+				actualMarkerFileBytes, errReadMarkerFile := os.ReadFile(actualMarkerFilePath)
+				require.NoError(t, errReadMarkerFile, "marker file should be readable")
+
+				expectedUpdateMarker := &UpdateMarker{
+					Version:           release.VersionWithSnapshot(),
+					Hash:              release.Commit(),
+					VersionedHome:     filepath.Join("data", fmt.Sprintf("elastic-agent-%s", release.ShortCommit())),
+					UpdatedOn:         aMomentInTime,
+					PrevVersion:       "1.2.3",
+					PrevHash:          "oldver",
+					PrevVersionedHome: "data/elastic-agent-1.2.3-oldver",
+					Details: &details.Details{
+						TargetVersion: release.VersionWithSnapshot(),
+						State:         details.StateRequested,
+					},
+					RollbacksAvailable: nil,
+				}
+
+				expectedMarkerBytes, err := yaml.Marshal(newMarkerSerializer(expectedUpdateMarker))
+				require.NoError(t, err, "error marshalling expected update marker")
+				require.YAMLEq(t, string(expectedMarkerBytes), string(actualMarkerFileBytes))
+			},
+		},
+		{
+			name: "no update marker, available install for rollback with expired TTL",
+			setup: func(t *testing.T, topDir string, agent *info.Agent, watcherHelper *MockWatcherHelper, installSource *mockInstallDescriptorSource) {
+				installSource.EXPECT().GetInstallDesc().Return(
+					&v1.InstallDescriptor{
+						AgentInstalls: []v1.AgentInstallDesc{
+							{
+								Version: release.VersionWithSnapshot(),
+								Hash:    release.Commit(),
+								// Versioned home should contain the version but since the path does not really exist we fallback to the legacy format with just the hash
+								//VersionedHome: filepath.Join("data", fmt.Sprintf("elastic-agent-%s-%s", release.VersionWithSnapshot(), release.ShortCommit())),
+								VersionedHome: filepath.Join("data", fmt.Sprintf("elastic-agent-%s", release.ShortCommit())),
+								Flavor:        "basic",
+								Active:        true,
+							},
+							{
+								// old install expired a second ago
+								OptionalTTLItem: v1.OptionalTTLItem{TTL: &aMomentAgo},
+								Version:         "1.2.3",
+								Hash:            "oldver",
+								VersionedHome:   "data/elastic-agent-1.2.3-oldver",
+								Flavor:          "basic",
+								Active:          false,
+							},
+						},
+					},
+					nil,
+				)
+			},
+			artifactSettings: artifact.DefaultConfig(),
+			upgradeSettings: &configuration.UpgradeConfig{
+				Rollback: &configuration.UpgradeRollbackConfig{
+					Window: 24 * time.Hour,
+				},
+			},
+			now:     aMomentInTime,
+			version: "1.2.3",
+			wantErr: func(t assert.TestingT, err error, i ...interface{}) bool {
+				return assert.ErrorIs(t, err, ErrNoRollbacksAvailable, i...)
+			},
+			additionalAsserts: func(t *testing.T, topDir string) {
+				actualMarkerFilePath := filepath.Join(topDir, "data", markerFilename)
+				require.NoFileExists(t, actualMarkerFilePath, "marker file must not be created")
+
+			},
+		},
 	}
 
 	for _, tc := range testcases {
@@ -281,7 +427,7 @@ func TestManualRollback(t *testing.T) {
 			require.NoError(t, err, "error creating data directory in topDir %q", topDir)
 
 			if tc.setup != nil {
-				tc.setup(t, topDir, mockAgentInfo, mockWatcherHelper)
+				tc.setup(t, topDir, mockAgentInfo, mockWatcherHelper, mockInstallSource)
 			}
 
 			upgrader, err := NewUpgrader(log, tc.artifactSettings, tc.upgradeSettings, mockAgentInfo, mockWatcherHelper, mockInstallSource)