From 2478442d0147651deedf2b0f376f2e0f3fd7ec70 Mon Sep 17 00:00:00 2001
From: Tengxiao Wang <126621870+twang-ps@users.noreply.github.com>
Date: Fri, 20 Sep 2024 14:55:21 -0700
Subject: [PATCH] Add SpeedUpForUT function to sched package (#2484)

* Add SpeedUpForUT function to sched package

---------

Signed-off-by: Tengxiao Wang <twang@purestorage.com>
---
 pkg/sched/intervals.go | 8 ++++++++
 1 file changed, 8 insertions(+)

diff --git a/pkg/sched/intervals.go b/pkg/sched/intervals.go
index a9353e787..c6acbd2b0 100644
--- a/pkg/sched/intervals.go
+++ b/pkg/sched/intervals.go
@@ -27,6 +27,8 @@ const (
 var (
 	// speedUp advances the clock faster for tests.
 	speedUp = false
+	// speedUpIntervalForUT overrides the interval of schedules in unit tests; no-op if set to 0
+	speedUpIntervalForUT = time.Duration(0)
 )
 
 // SpeedUp advances the clock faster for tests.
@@ -38,6 +40,10 @@ func inSpeedUp() bool {
 	return speedUp
 }
 
+func SpeedUpForUT(duration time.Duration) {
+	speedUpIntervalForUT = duration
+}
+
 type IntervalSpec struct {
 	Freq    string
 	Period  uint64 `yaml:"period,omitempty"`
@@ -560,6 +566,8 @@ func (p RetainIntervalImpl) nextAfter(t time.Time) time.Time {
 	newTime := p.iv.nextAfter(t)
 	if inSpeedUp() {
 		return t.Add(time.Minute)
+	} else if speedUpIntervalForUT != time.Duration(0) {
+		return t.Add(speedUpIntervalForUT)
 	}
 	return newTime
 }