From 05a82a43dc354df2cc4746a8cdc0db1e8a5e3107 Mon Sep 17 00:00:00 2001
From: Kegan Dougal <7190048+kegsay@users.noreply.github.com>
Date: Mon, 11 Mar 2024 12:06:13 +0000
Subject: [PATCH] Same race pattern as timeSince for timeSleep

---
 sync2/poller_test.go | 47 ++++++++++++++++++++++++++++++--------------
 1 file changed, 32 insertions(+), 15 deletions(-)

diff --git a/sync2/poller_test.go b/sync2/poller_test.go
index 1ce7a74..4b145aa 100644
--- a/sync2/poller_test.go
+++ b/sync2/poller_test.go
@@ -25,12 +25,23 @@ const initialSinceToken = "0"
 var (
 	timeSinceMu    sync.Mutex
 	timeSinceValue = time.Duration(0) // 0 means use the real impl
+	timeSleepMu    sync.Mutex
+	timeSleepValue = time.Duration(0)  // 0 means use the real impl
+	timeSleepCheck func(time.Duration) // called to check sleep values
 )
 
 func setTimeSinceValue(val time.Duration) {
 	timeSinceMu.Lock()
+	defer timeSinceMu.Unlock()
 	timeSinceValue = val
-	timeSinceMu.Unlock()
+}
+func setTimeSleepDelay(val time.Duration, fn ...func(d time.Duration)) {
+	timeSleepMu.Lock()
+	defer timeSleepMu.Unlock()
+	timeSleepValue = val
+	if len(fn) > 0 {
+		timeSleepCheck = fn[0]
+	}
 }
 func init() {
 	timeSince = func(t time.Time) time.Duration {
@@ -41,6 +52,18 @@ func init() {
 		}
 		return timeSinceValue
 	}
+	timeSleep = func(d time.Duration) {
+		timeSleepMu.Lock()
+		defer timeSleepMu.Unlock()
+		if timeSleepCheck != nil {
+			timeSleepCheck(d)
+		}
+		if timeSleepValue == 0 {
+			time.Sleep(d)
+			return
+		}
+		time.Sleep(timeSleepValue)
+	}
 }
 
 // Tests that EnsurePolling works in the happy case
@@ -583,12 +606,10 @@ func TestPollerGivesUpEventually(t *testing.T) {
 	accumulator, client := newMocks(func(authHeader, since string) (*SyncResponse, int, error) {
 		return nil, 524, fmt.Errorf("gateway timeout")
 	})
-	timeSleep = func(d time.Duration) {
-		// actually sleep to make sure async actions can happen if any
-		time.Sleep(1 * time.Microsecond)
-	}
+	// actually sleep to make sure async actions can happen if any
+	setTimeSleepDelay(time.Microsecond)
 	defer func() { // reset the value after the test runs
-		timeSleep = time.Sleep
+		setTimeSleepDelay(0)
 	}()
 	var wg sync.WaitGroup
 	wg.Add(1)
@@ -654,15 +675,13 @@ func TestPollerBackoff(t *testing.T) {
 		wantBackoffDuration = errorResponses[i].backoff
 		return nil, errorResponses[i].code, errorResponses[i].err
 	})
-	timeSleep = func(d time.Duration) {
+	setTimeSleepDelay(time.Millisecond, func(d time.Duration) {
 		if d != wantBackoffDuration {
 			t.Errorf("time.Sleep called incorrectly: got %v want %v", d, wantBackoffDuration)
 		}
-		// actually sleep to make sure async actions can happen if any
-		time.Sleep(1 * time.Millisecond)
-	}
+	})
 	defer func() { // reset the value after the test runs
-		timeSleep = time.Sleep
+		setTimeSleepDelay(0)
 	}()
 	var wg sync.WaitGroup
 	wg.Add(1)
@@ -727,12 +746,10 @@ func TestPollerResendsOnCallbackError(t *testing.T) {
 	pid := PollerID{UserID: "@TestPollerResendsOnCallbackError:localhost", DeviceID: "FOOBAR"}
 
 	defer func() { // reset the value after the test runs
-		timeSleep = time.Sleep
+		setTimeSleepDelay(0)
 	}()
 	// we don't actually want to wait 3s between retries, so monkey patch it out
-	timeSleep = func(d time.Duration) {
-		time.Sleep(time.Millisecond)
-	}
+	setTimeSleepDelay(time.Millisecond)
 
 	testCases := []struct {
 		name             string