Wait for upgrade details

testing/upgradetest/upgrader.go

@@ -254,8 +254,8 @@ func PerformUpgrade(
 
 	// Check that, while the Upgrade Watcher is running, the upgrade details in Agent status
 	// show the state as UPG_WATCHING.
-	if err := checkUpgradeDetailsState(ctx, startFixture, details.StateWatching); err != nil {
-		// error context added by checkUpgradeDetailsState
+	if err := waitUpgradeDetailsState(ctx, startFixture, details.StateWatching, 2*time.Minute, 10*time.Second); err != nil {
+		// error context added by waitUpgradeDetailsState
 		return err
 	}
 
@@ -300,8 +300,8 @@ func PerformUpgrade(
 
 	// Check that, upon successful upgrade, the upgrade details have been cleared out
 	// from Agent status.
-	if err := checkUpgradeDetailsState(ctx, startFixture, ""); err != nil {
-		// error context added by checkUpgradeDetailsState
+	if err := waitUpgradeDetailsState(ctx, startFixture, "", 2*time.Minute, 10*time.Second); err != nil {
+		// error context added by waitUpgradeDetailsState
 		return err
 	}
 
@@ -403,29 +403,50 @@ func WaitHealthyAndVersion(ctx context.Context, f *atesting.Fixture, versionInfo
 	}
 }
 
-func checkUpgradeDetailsState(ctx context.Context, f *atesting.Fixture, expectedState details.State) error {
-	status, err := f.ExecStatus(ctx)
-	if err != nil {
-		return err
-	}
+func waitUpgradeDetailsState(ctx context.Context, f *atesting.Fixture, expectedState details.State, timeout time.Duration, interval time.Duration) error {
+	ctx, cancel := context.WithTimeout(ctx, timeout)
+	defer cancel()
 
-	if expectedState == "" {
-		if status.UpgradeDetails == nil {
-			// Expected and actual match, so we're good
-			return nil
-		}
+	t := time.NewTicker(interval)
+	defer t.Stop()
 
-		return errors.New("upgrade details found in status but they were expected to be absent")
-	}
+	var lastErr error
+	for {
+		select {
+		case <-ctx.Done():
+			if lastErr != nil {
+				return fmt.Errorf("failed waiting for healthy agent and version (%w): %w", ctx.Err(), lastErr)
+			}
+			return ctx.Err()
+		case <-t.C:
+			status, err := f.ExecStatus(ctx)
+			if err != nil {
+				lastErr = err
+				continue
+			}
 
-	if status.UpgradeDetails == nil {
-		return fmt.Errorf("upgrade details not found in status but expected upgrade details state was [%s]", expectedState)
-	}
+			if expectedState == "" {
+				if status.UpgradeDetails == nil {
+					// Expected and actual match, so we're good
+					return nil
+				}
 
-	// Neither expected nor actual are nil, so compare the two
-	if status.UpgradeDetails.State == expectedState {
-		return nil
-	}
+				lastErr = errors.New("upgrade details found in status but they were expected to be absent")
+				continue
+			}
+
+			if status.UpgradeDetails == nil {
+				lastErr = fmt.Errorf("upgrade details not found in status but expected upgrade details state was [%s]", expectedState)
+				continue
+			}
+
+			// Neither expected nor actual are nil, so compare the two
+			if status.UpgradeDetails.State == expectedState {
+				return nil
+			}
 
-	return fmt.Errorf("upgrade details state in status [%s] is not the same as expected upgrade details state  [%s]", status.UpgradeDetails.State, expectedState)
+			lastErr = fmt.Errorf("upgrade details state in status [%s] is not the same as expected upgrade details state  [%s]", status.UpgradeDetails.State, expectedState)
+			continue
+		}
+	}
 }