Merge pull request #3 from tkuchiki/patch-86

patch jrallison#86

Author: tkuchiki

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"`

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)
 		})
 	})