Fix TestStandaloneUpgradeRollback (#5445)

This commit fixes two problems:
1. When getting the Elastic-Agent status to check for
state.UpgradeDetails, it could happen that the upgrade had already
finished, thus state.UpgradeDetails was nil. Now we only check for
state.UpgradeDetails if the status is upgrading
2. When CheckHealthyAndVersion checks for the version during the
rollback it could get the previous version, thus failing. This is
fixed by ignoring version mismatch error.

testing/integration/upgrade_rollback_test.go

@@ -23,6 +23,7 @@ import (
 	"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/install"
+	"github.com/elastic/elastic-agent/pkg/control/v2/cproto"
 	atesting "github.com/elastic/elastic-agent/pkg/testing"
 	"github.com/elastic/elastic-agent/pkg/testing/define"
 	"github.com/elastic/elastic-agent/pkg/testing/tools/testcontext"
@@ -135,8 +136,14 @@ inputs:
 		state, err := client.State(ctx)
 		require.NoError(t, err)
 
-		require.NotNil(t, state.UpgradeDetails, "upgrade details in the state cannot be nil")
-		require.Equal(t, details.StateRollback, details.State(state.UpgradeDetails.State))
+		if state.State == cproto.State_UPGRADING {
+			if state.UpgradeDetails == nil {
+				t.Fatal("upgrade details in the state cannot be nil")
+			}
+			require.Equal(t, details.StateRollback, details.State(state.UpgradeDetails.State))
+		} else {
+			t.Logf("rollback finished, status is '%s', cannot check UpgradeDetails", state.State.String())
+		}
 	}
 
 	// rollback should stop the watcher

testing/upgradetest/upgrader.go

@@ -421,14 +421,17 @@ func PerformUpgrade(
 	return nil
 }
 
+var ErrVerMismatch = errors.New("versions don't match")
+
 func CheckHealthyAndVersion(ctx context.Context, f *atesting.Fixture, versionInfo atesting.AgentBinaryVersion) error {
 	checkFunc := func() error {
 		status, err := f.ExecStatus(ctx)
 		if err != nil {
 			return err
 		}
 		if status.Info.Version != versionInfo.Version {
-			return fmt.Errorf("versions don't match: got %s, want %s",
+			return fmt.Errorf("%w: got %s, want %s",
+				ErrVerMismatch,
 				status.Info.Version, versionInfo.Version)
 		}
 		if status.Info.Snapshot != versionInfo.Snapshot {
@@ -494,6 +497,9 @@ func WaitHealthyAndVersion(ctx context.Context, f *atesting.Fixture, versionInfo
 	ctx, cancel := context.WithTimeout(ctx, timeout)
 	defer cancel()
 
+	// The deadline was set above, we don't need to check for it.
+	deadline, _ := ctx.Deadline()
+
 	t := time.NewTicker(interval)
 	defer t.Stop()
 
@@ -507,6 +513,12 @@ func WaitHealthyAndVersion(ctx context.Context, f *atesting.Fixture, versionInfo
 			return ctx.Err()
 		case <-t.C:
 			err := CheckHealthyAndVersion(ctx, f, versionInfo)
+			// If we're in an upgrade process, the versions might not match
+			// so we wait to see if we get to a stable version
+			if errors.Is(err, ErrVerMismatch) {
+				logger.Logf("version mismatch, ignoring it, time until timeout: %s", deadline.Sub(time.Now()))
+				continue
+			}
 			if err == nil {
 				return nil
 			}