test: do not wait afterRestartDelay in TestCleanup (#5446)

delay after agent restart is performed to allow agent to tear down all the processes
important mainly for windows, as it prevents removing files which are in use.
In TestCleanup there are no processes and we don't need to wait.
Update Cleanup method to accept delay as an arg and defaulting to afterRestartDelay in prod

TestCleanup speed improves depending on the os:

linux: 8s -> 0.2s
darwin 8s -> 0.2s
windows: 60s -> 0.2s

Author: kruskall

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

@@ -68,8 +68,12 @@ func Rollback(ctx context.Context, log *logger.Logger, c client.Client, topDirPa
 
 // Cleanup removes all artifacts and files related to a specified version.
 func Cleanup(log *logger.Logger, topDirPath, currentVersionedHome, currentHash string, removeMarker, keepLogs bool) error {
+	return cleanup(log, topDirPath, currentVersionedHome, currentHash, removeMarker, keepLogs, afterRestartDelay)
+}
+
+func cleanup(log *logger.Logger, topDirPath, currentVersionedHome, currentHash string, removeMarker, keepLogs bool, delay time.Duration) error {
 	log.Infow("Cleaning up upgrade", "hash", currentHash, "remove_marker", removeMarker)
-	<-time.After(afterRestartDelay)
+	<-time.After(delay)
 
 	// data directory path
 	dataDirPath := paths.DataFrom(topDirPath)

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

@@ -215,7 +215,7 @@ func TestCleanup(t *testing.T) {
 			require.NoError(t, err, "error loading update marker")
 			require.NotNil(t, marker, "loaded marker must not be nil")
 			t.Logf("Loaded update marker %+v", marker)
-			tt.wantErr(t, Cleanup(testLogger, testTop, marker.VersionedHome, marker.Hash, tt.args.removeMarker, tt.args.keepLogs), fmt.Sprintf("Cleanup(%v, %v, %v, %v)", marker.VersionedHome, marker.Hash, tt.args.removeMarker, tt.args.keepLogs))
+			tt.wantErr(t, cleanup(testLogger, testTop, marker.VersionedHome, marker.Hash, tt.args.removeMarker, tt.args.keepLogs, 0), fmt.Sprintf("Cleanup(%v, %v, %v, %v)", marker.VersionedHome, marker.Hash, tt.args.removeMarker, tt.args.keepLogs))
 			tt.checkAfterCleanup(t, testTop)
 		})
 	}