backport subtle test fix from #16396

hashicorp · May 10, 2023 · 921c184 · 921c184
1 parent 0997ddf
commit 921c184
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 1 deletion.
diff --git a/agent/grpc-external/services/peerstream/stream_test.go b/agent/grpc-external/services/peerstream/stream_test.go
@@ -1161,8 +1161,8 @@ func TestStreamResources_Server_DisconnectsOnHeartbeatTimeout(t *testing.T) {
 	})
 
 	testutil.RunStep(t, "stream is disconnected due to heartbeat timeout", func(t *testing.T) {
-		disconnectTime := it.FutureNow(1)
 		retry.Run(t, func(r *retry.R) {
+			disconnectTime := it.StaticNow()
 			status, ok := srv.StreamStatus(testPeerID)
 			require.True(r, ok)
 			require.False(r, status.Connected)

diff --git a/agent/grpc-external/services/peerstream/testing.go b/agent/grpc-external/services/peerstream/testing.go
@@ -130,6 +130,16 @@ func (t *incrementalTime) Now() time.Time {
 	return t.base.Add(dur)
 }
 
+// StaticNow returns the current internal clock without advancing it.
+func (t *incrementalTime) StaticNow() time.Time {
+	t.mu.Lock()
+	defer t.mu.Unlock()
+
+	dur := time.Duration(t.next) * time.Second
+
+	return t.base.Add(dur)
+}
+
 // FutureNow will return a given future value of the Now() function.
 // The numerical argument indicates which future Now value you wanted. The
 // value must be > 0.