TEST: fix redis flaky tests; take 2 (#109)

taylorchu · Jul 5, 2023 · 733c924 · 733c924
1 parent cf9646e
commit 733c924
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 21 deletions.
diff --git a/middleware/heartbeat/heartbeater.go b/middleware/heartbeat/heartbeater.go
@@ -23,9 +23,9 @@ type HeartbeaterOptions struct {
 func Heartbeater(hopts *HeartbeaterOptions) work.HandleMiddleware {
 	return func(f work.HandleFunc) work.HandleFunc {
 		return func(job *work.Job, opt *work.DequeueOptions) error {
+			copiedJob := *job
 			refresh := func() error {
 				now := time.Now()
-				copiedJob := *job
 				copiedJob.UpdatedAt = now
 				copiedJob.EnqueuedAt = now.Add(time.Duration(hopts.InvisibleSec) * time.Second)
 				return hopts.Queue.Enqueue(&copiedJob, &work.EnqueueOptions{
@@ -57,6 +57,8 @@ func Heartbeater(hopts *HeartbeaterOptions) work.HandleMiddleware {
 			err := f(job, opt)
 			cancel()
 			<-done
+			job.UpdatedAt = copiedJob.UpdatedAt
+			job.EnqueuedAt = copiedJob.EnqueuedAt
 			return err
 		}
 	}

diff --git a/middleware/heartbeat/heartbeater_test.go b/middleware/heartbeat/heartbeater_test.go
@@ -27,7 +27,7 @@ func TestHeartbeater(t *testing.T) {
 	hb := Heartbeater(&HeartbeaterOptions{
 		Queue:        work.NewRedisQueue(client),
 		InvisibleSec: 30,
-		IntervalSec:  10,
+		IntervalSec:  1,
 	})
 
 	h := hb(func(*work.Job, *work.DequeueOptions) error {
@@ -36,8 +36,8 @@ func TestHeartbeater(t *testing.T) {
 
 	err := h(job, opt)
 	require.NoError(t, err)
-	require.Equal(t, job.CreatedAt, job.UpdatedAt)
-	require.Equal(t, job.CreatedAt, job.EnqueuedAt)
+	require.Equal(t, job.EnqueuedAt.Unix(), job.UpdatedAt.Unix()+30)
+	require.NotEqual(t, job.CreatedAt, job.UpdatedAt)
 
 	z, err := client.ZRangeByScoreWithScores(
 		context.Background(),
@@ -48,21 +48,5 @@ func TestHeartbeater(t *testing.T) {
 		}).Result()
 	require.NoError(t, err)
 	require.Len(t, z, 1)
-	require.EqualValues(t, job.EnqueuedAt.Unix()+30, z[0].Score)
-
-	err = h(job, opt)
-	require.NoError(t, err)
-	require.Equal(t, job.CreatedAt, job.UpdatedAt)
-	require.Equal(t, job.CreatedAt, job.EnqueuedAt)
-
-	z, err = client.ZRangeByScoreWithScores(
-		context.Background(),
-		"{ns1}:queue:q1",
-		&redis.ZRangeBy{
-			Min: "-inf",
-			Max: "+inf",
-		}).Result()
-	require.NoError(t, err)
-	require.Len(t, z, 1)
-	require.EqualValues(t, job.EnqueuedAt.Unix()+30, z[0].Score)
+	require.EqualValues(t, job.EnqueuedAt.Unix(), z[0].Score)
 }