Add KillRetryTask, KillScheduledTask methods to RDB

internal/rdb/inspect_test.go

@@ -832,6 +832,188 @@ func TestEnqueueAllDeadTasks(t *testing.T) {
 	}
 }
 
+func TestKillRetryTask(t *testing.T) {
+	r := setup(t)
+	m1 := newTaskMessage("send_email", nil)
+	m2 := newTaskMessage("reindex", nil)
+	t1 := time.Now().Add(time.Minute)
+	t2 := time.Now().Add(time.Hour)
+
+	tests := []struct {
+		retry     []sortedSetEntry
+		dead      []sortedSetEntry
+		id        xid.ID
+		score     int64
+		want      error
+		wantRetry []sortedSetEntry
+		wantDead  []sortedSetEntry
+	}{
+		{
+			retry: []sortedSetEntry{
+				{m1, t1.Unix()},
+				{m2, t2.Unix()},
+			},
+			dead:  []sortedSetEntry{},
+			id:    m1.ID,
+			score: t1.Unix(),
+			want:  nil,
+			wantRetry: []sortedSetEntry{
+				{m2, t2.Unix()},
+			},
+			wantDead: []sortedSetEntry{
+				{m1, time.Now().Unix()},
+			},
+		},
+		{
+			retry: []sortedSetEntry{
+				{m1, t1.Unix()},
+			},
+			dead: []sortedSetEntry{
+				{m2, t2.Unix()},
+			},
+			id:    m2.ID,
+			score: t2.Unix(),
+			want:  ErrTaskNotFound,
+			wantRetry: []sortedSetEntry{
+				{m1, t1.Unix()},
+			},
+			wantDead: []sortedSetEntry{
+				{m2, t2.Unix()},
+			},
+		},
+	}
+
+	for _, tc := range tests {
+		flushDB(t, r)
+		seedRetryQueue(t, r, tc.retry)
+		seedDeadQueue(t, r, tc.dead)
+
+		got := r.KillRetryTask(tc.id, tc.score)
+		if got != tc.want {
+			t.Errorf("(*RDB).KillRetryTask(%v, %v) = %v, want %v",
+				tc.id, tc.score, got, tc.want)
+			continue
+		}
+
+		gotRetryRaw := r.client.ZRangeWithScores(base.RetryQueue, 0, -1).Val()
+		var gotRetry []sortedSetEntry
+		for _, z := range gotRetryRaw {
+			gotRetry = append(gotRetry, sortedSetEntry{
+				Msg:   mustUnmarshal(t, z.Member.(string)),
+				Score: int64(z.Score),
+			})
+		}
+		if diff := cmp.Diff(tc.wantRetry, gotRetry, sortZSetEntryOpt, timeCmpOpt); diff != "" {
+			t.Errorf("mismatch found in %q; (-want,+got)\n%s",
+				base.RetryQueue, diff)
+		}
+
+		gotDeadRaw := r.client.ZRangeWithScores(base.DeadQueue, 0, -1).Val()
+		var gotDead []sortedSetEntry
+		for _, z := range gotDeadRaw {
+			gotDead = append(gotDead, sortedSetEntry{
+				Msg:   mustUnmarshal(t, z.Member.(string)),
+				Score: int64(z.Score),
+			})
+		}
+		if diff := cmp.Diff(tc.wantDead, gotDead, sortZSetEntryOpt, timeCmpOpt); diff != "" {
+			t.Errorf("mismatch found in %q; (-want,+got)\n%s",
+				base.DeadQueue, diff)
+		}
+	}
+}
+
+func TestKillScheduledTask(t *testing.T) {
+	r := setup(t)
+	m1 := newTaskMessage("send_email", nil)
+	m2 := newTaskMessage("reindex", nil)
+	t1 := time.Now().Add(time.Minute)
+	t2 := time.Now().Add(time.Hour)
+
+	tests := []struct {
+		scheduled     []sortedSetEntry
+		dead          []sortedSetEntry
+		id            xid.ID
+		score         int64
+		want          error
+		wantScheduled []sortedSetEntry
+		wantDead      []sortedSetEntry
+	}{
+		{
+			scheduled: []sortedSetEntry{
+				{m1, t1.Unix()},
+				{m2, t2.Unix()},
+			},
+			dead:  []sortedSetEntry{},
+			id:    m1.ID,
+			score: t1.Unix(),
+			want:  nil,
+			wantScheduled: []sortedSetEntry{
+				{m2, t2.Unix()},
+			},
+			wantDead: []sortedSetEntry{
+				{m1, time.Now().Unix()},
+			},
+		},
+		{
+			scheduled: []sortedSetEntry{
+				{m1, t1.Unix()},
+			},
+			dead: []sortedSetEntry{
+				{m2, t2.Unix()},
+			},
+			id:    m2.ID,
+			score: t2.Unix(),
+			want:  ErrTaskNotFound,
+			wantScheduled: []sortedSetEntry{
+				{m1, t1.Unix()},
+			},
+			wantDead: []sortedSetEntry{
+				{m2, t2.Unix()},
+			},
+		},
+	}
+
+	for _, tc := range tests {
+		flushDB(t, r)
+		seedScheduledQueue(t, r, tc.scheduled)
+		seedDeadQueue(t, r, tc.dead)
+
+		got := r.KillScheduledTask(tc.id, tc.score)
+		if got != tc.want {
+			t.Errorf("(*RDB).KillScheduledTask(%v, %v) = %v, want %v",
+				tc.id, tc.score, got, tc.want)
+			continue
+		}
+
+		gotScheduledRaw := r.client.ZRangeWithScores(base.ScheduledQueue, 0, -1).Val()
+		var gotScheduled []sortedSetEntry
+		for _, z := range gotScheduledRaw {
+			gotScheduled = append(gotScheduled, sortedSetEntry{
+				Msg:   mustUnmarshal(t, z.Member.(string)),
+				Score: int64(z.Score),
+			})
+		}
+		if diff := cmp.Diff(tc.wantScheduled, gotScheduled, sortZSetEntryOpt, timeCmpOpt); diff != "" {
+			t.Errorf("mismatch found in %q; (-want,+got)\n%s",
+				base.ScheduledQueue, diff)
+		}
+
+		gotDeadRaw := r.client.ZRangeWithScores(base.DeadQueue, 0, -1).Val()
+		var gotDead []sortedSetEntry
+		for _, z := range gotDeadRaw {
+			gotDead = append(gotDead, sortedSetEntry{
+				Msg:   mustUnmarshal(t, z.Member.(string)),
+				Score: int64(z.Score),
+			})
+		}
+		if diff := cmp.Diff(tc.wantDead, gotDead, sortZSetEntryOpt, timeCmpOpt); diff != "" {
+			t.Errorf("mismatch found in %q; (-want,+got)\n%s",
+				base.DeadQueue, diff)
+		}
+	}
+}
+
 func TestDeleteDeadTask(t *testing.T) {
 	r := setup(t)
 	m1 := newTaskMessage("send_email", nil)