From c52b042b8cfa135541feba3d094e7cae7902fb5b Mon Sep 17 00:00:00 2001
From: tkuchiki <kuchiki.taku@gmail.com>
Date: Mon, 24 Aug 2020 21:56:43 +0900
Subject: [PATCH] patch #86

---
 enqueue.go      | 22 ++++++++++++++++++++++
 enqueue_test.go | 22 +++++++++++++++++-----
 2 files changed, 39 insertions(+), 5 deletions(-)

diff --git a/enqueue.go b/enqueue.go
index f9af7a8..49b5d1a 100644
--- a/enqueue.go
+++ b/enqueue.go
@@ -21,6 +21,28 @@ type EnqueueData struct {
 	EnqueueOptions
 }
 
+type EnqueueDataProxy EnqueueData
+
+func (e EnqueueData) MarshalJSON() ([]byte, error) {
+	o := e.EnqueueOptions
+	if e.EnqueueOptions.RetryCount > 0 {
+		s := struct {
+			EnqueueDataProxy
+			Retry      int     `json:"retry,omitempty"`
+			RetryCount int     `json:"retry_count"`
+			At         float64 `json:"at,omitempty"`
+		}{EnqueueDataProxy(e), o.RetryCount, 0, o.At}
+		return json.Marshal(s)
+	}
+
+	return json.Marshal(struct {
+		EnqueueDataProxy
+		RetryCount int     `json:"retry_count,omitempty"`
+		Retry      bool    `json:"retry,omitempty"`
+		At         float64 `json:"at,omitempty"`
+	}{EnqueueDataProxy(e), o.RetryCount, o.Retry, o.At})
+}
+
 type EnqueueOptions struct {
 	RetryCount int     `json:"retry_count,omitempty"`
 	Retry      bool    `json:"retry,omitempty"`
diff --git a/enqueue_test.go b/enqueue_test.go
index 3980941..9a8caa5 100644
--- a/enqueue_test.go
+++ b/enqueue_test.go
@@ -72,8 +72,23 @@ func EnqueueSpec(c gospec.Context) {
 			c.Expect(ea, IsWithin(0.1), nowToSecondsWithNanoPrecision())
 		})
 
-		c.Specify("has retry and retry_count when set", func() {
-			EnqueueWithOptions("enqueue6", "Compare", []string{"foo", "bar"}, EnqueueOptions{RetryCount: 13, Retry: true})
+		c.Specify("sets retry count to `retry`", func() {
+			EnqueueWithOptions("enqueue6", "Compare", []string{"foo", "bar"}, EnqueueOptions{RetryCount: 13})
+
+			bytes, _ := redis.Bytes(conn.Do("lpop", "prod:queue:enqueue6"))
+			var result map[string]interface{}
+			json.Unmarshal(bytes, &result)
+			c.Expect(result["class"], Equals, "Compare")
+
+			retry := result["retry"].(float64)
+			c.Expect(retry, Equals, float64(13))
+
+			retryCount := result["retry_count"].(float64)
+			c.Expect(retryCount, Equals, float64(0))
+		})
+
+		c.Specify("sets Retry correctly when no count given", func() {
+			EnqueueWithOptions("enqueue6", "Compare", []string{"foo", "bar"}, EnqueueOptions{Retry: true})
 
 			bytes, _ := redis.Bytes(conn.Do("lpop", "prod:queue:enqueue6"))
 			var result map[string]interface{}
@@ -82,9 +97,6 @@ func EnqueueSpec(c gospec.Context) {
 
 			retry := result["retry"].(bool)
 			c.Expect(retry, Equals, true)
-
-			retryCount := int(result["retry_count"].(float64))
-			c.Expect(retryCount, Equals, 13)
 		})
 	})