From bdbee31b861381fe6285034d6af8a6167ce82b66 Mon Sep 17 00:00:00 2001
From: Ken Hibino <ken.hibino7@gmail.com>
Date: Thu, 28 Nov 2019 11:22:42 -0800
Subject: [PATCH] Add test for retry task logic

---
 retry.go      | 12 ++++-----
 retry_test.go | 69 +++++++++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 74 insertions(+), 7 deletions(-)
 create mode 100644 retry_test.go

diff --git a/retry.go b/retry.go
index 6e93c07..b8fda87 100644
--- a/retry.go
+++ b/retry.go
@@ -1,7 +1,6 @@
 package asynq
 
 import (
-	"fmt"
 	"log"
 	"math"
 	"math/rand"
@@ -9,20 +8,19 @@ import (
 )
 
 func retryTask(rdb *rdb, msg *taskMessage, err error) {
+	msg.ErrorMsg = err.Error()
 	if msg.Retried >= msg.Retry {
-		fmt.Println("[DEBUG] Retry exhausted!!!")
+		log.Printf("[WARN] Retry exhausted for task(Type: %q, ID: %v)\n", msg.Type, msg.ID)
 		if err := rdb.kill(msg); err != nil {
-			log.Printf("[ERROR] could not add task %+v to 'dead' set\n", err)
+			log.Printf("[ERROR] Could not add task %+v to 'dead'\n", err)
 		}
 		return
 	}
 	retryAt := time.Now().Add(delaySeconds((msg.Retried)))
-	fmt.Printf("[DEBUG] Retrying the task in %v\n", retryAt.Sub(time.Now()))
+	log.Printf("[INFO] Retrying task(Type: %q, ID: %v) in %v\n", msg.Type, msg.ID, retryAt.Sub(time.Now()))
 	msg.Retried++
-	msg.ErrorMsg = err.Error()
 	if err := rdb.schedule(retry, retryAt, msg); err != nil {
-		// TODO(hibiken): Not sure how to handle this error
-		log.Printf("[ERROR] could not add msg %+v to 'retry' set: %v\n", msg, err)
+		log.Printf("[ERROR] Could not add msg %+v to 'retry': %v\n", msg, err)
 		return
 	}
 }
diff --git a/retry_test.go b/retry_test.go
new file mode 100644
index 0000000..600136f
--- /dev/null
+++ b/retry_test.go
@@ -0,0 +1,69 @@
+package asynq
+
+import (
+	"fmt"
+	"testing"
+
+	"github.com/google/go-cmp/cmp"
+	"github.com/google/uuid"
+)
+
+func TestRetry(t *testing.T) {
+	r := setup(t)
+	errMsg := "email server not responding"
+	// t1 is a task with max-retry count reached.
+	t1 := &taskMessage{Type: "send_email", Retry: 10, Retried: 10, Queue: "default", ID: uuid.New()}
+	// t2 is t1 with updated error message.
+	t2 := *t1
+	t2.ErrorMsg = errMsg
+	// t3 is a task which hasn't reached max-retry count.
+	t3 := &taskMessage{Type: "send_email", Retry: 10, Retried: 5, Queue: "default", ID: uuid.New()}
+	// t4 is t3 after retry.
+	t4 := *t3
+	t4.Retried++
+	t4.ErrorMsg = errMsg
+
+	tests := []struct {
+		desc      string         // test case description
+		msg       *taskMessage   // task to retry
+		err       error          // error that caused retry
+		wantDead  []*taskMessage // state "dead" queue should be in
+		wantRetry []*taskMessage // state "retry" queue should be in
+	}{
+		{
+			desc:      "With retry exhausted task",
+			msg:       t1,
+			err:       fmt.Errorf(errMsg),
+			wantDead:  []*taskMessage{&t2},
+			wantRetry: []*taskMessage{},
+		},
+		{
+			desc:      "With retry-able task",
+			msg:       t3,
+			err:       fmt.Errorf(errMsg),
+			wantDead:  []*taskMessage{},
+			wantRetry: []*taskMessage{&t4},
+		},
+	}
+
+	for _, tc := range tests {
+		// clean up db before each test case.
+		if err := r.client.FlushDB().Err(); err != nil {
+			t.Fatal(err)
+		}
+
+		retryTask(r, tc.msg, tc.err)
+
+		deadQueue := r.client.ZRange(dead, 0, -1).Val()
+		gotDead := mustUnmarshalSlice(t, deadQueue)
+		if diff := cmp.Diff(tc.wantDead, gotDead, sortMsgOpt); diff != "" {
+			t.Errorf("%s;\nmismatch found in %q after retryTask(); (-want, +got)\n%s", tc.desc, dead, diff)
+		}
+
+		retryQueue := r.client.ZRange(retry, 0, -1).Val()
+		gotRetry := mustUnmarshalSlice(t, retryQueue)
+		if diff := cmp.Diff(tc.wantRetry, gotRetry, sortMsgOpt); diff != "" {
+			t.Errorf("%s;\nmismatch found in %q after retryTask(); (-want, +got)\n%s", tc.desc, dead, diff)
+		}
+	}
+}