diff --git a/internal/pkg/agent/application/upgrade/rollback.go b/internal/pkg/agent/application/upgrade/rollback.go
index b16708d90b7..81f3d80c0e8 100644
--- a/internal/pkg/agent/application/upgrade/rollback.go
+++ b/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)
diff --git a/internal/pkg/agent/application/upgrade/rollback_test.go b/internal/pkg/agent/application/upgrade/rollback_test.go
index e87e2dc0beb..f85ef5c4607 100644
--- a/internal/pkg/agent/application/upgrade/rollback_test.go
+++ b/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)
 		})
 	}