Add postWatcherSuccessHook to upgrade test

Author: ycombinator

testing/upgradetest/upgrader.go

@@ -51,10 +51,11 @@ type upgradeOpts struct {
 	// Disable check that enforces different hashed between the to and from version of upgrade
 	disableHashCheck bool
 
-	preInstallHook  func() error
-	postInstallHook func() error
-	preUpgradeHook  func() error
-	postUpgradeHook func() error
+	preInstallHook         func() error
+	postInstallHook        func() error
+	preUpgradeHook         func() error
+	postUpgradeHook        func() error
+	postWatcherSuccessHook func(context.Context, *atesting.Fixture) error
 }
 
 type UpgradeOpt func(opts *upgradeOpts)
@@ -125,6 +126,14 @@ func WithPostUpgradeHook(hook func() error) UpgradeOpt {
 	}
 }
 
+// WithPostWatcherSuccessHook sets a hook to be called after the upgrade is successful
+// and the upgrade watcher has terminated as well.
+func WithPostWatcherSuccessHook(hook func(context.Context, *atesting.Fixture) error) UpgradeOpt {
+	return func(opts *upgradeOpts) {
+		opts.postWatcherSuccessHook = hook
+	}
+}
+
 // WithCustomWatcherConfig sets a custom watcher configuration to use.
 func WithCustomWatcherConfig(cfg string) UpgradeOpt {
 	return func(opts *upgradeOpts) {
@@ -435,6 +444,12 @@ func PerformUpgrade(
 		}
 	}
 
+	if upgradeOpts.postWatcherSuccessHook != nil {
+		if err := upgradeOpts.postWatcherSuccessHook(ctx, endFixture); err != nil {
+			return fmt.Errorf("post watcher success hook failed: %w", err)
+		}
+	}
+
 	return nil
 }