Merge pull request #164 from mkuratczyk/issue-149

Add support for numeric delete-after values
michaelklishin · Sep 24, 2020 · 555e60e · 555e60e
2 parents 4af0caa + fb388c5
commit 555e60e
Show file tree

Hide file tree

Showing 2 changed files with 106 additions and 26 deletions.
diff --git a/rabbithole_test.go b/rabbithole_test.go
@@ -2422,11 +2422,57 @@ var _ = Describe("Rabbithole", func() {
 				DeleteAfter:       "never"}
 
 			_, err := rmqc.DeclareShovel(vh, sn, shovelDefinition)
-			Ω(err).Should(BeNil())
+			Ω(err).Should(BeNil(), "Error declaring shovel")
 
 			awaitEventPropagation()
 			x, err := rmqc.GetShovel(vh, sn)
+			Ω(err).Should(BeNil(), "Error getting shovel")
+			Ω(x.Name).Should(Equal(sn))
+			Ω(x.Vhost).Should(Equal(vh))
+			Ω(x.Component).Should(Equal("shovel"))
+			Ω(x.Definition.SourceURI).Should(Equal(ssu))
+			Ω(x.Definition.SourceQueue).Should(Equal("mySourceQueue"))
+			Ω(x.Definition.DestinationURI).Should(Equal(sdu))
+			Ω(x.Definition.DestinationQueue).Should(Equal("myDestQueue"))
+			Ω(x.Definition.AddForwardHeaders).Should(Equal(true))
+			Ω(x.Definition.AckMode).Should(Equal("on-confirm"))
+			Ω(string(x.Definition.DeleteAfter)).Should(Equal("never"))
+
+			_, err = rmqc.DeleteShovel(vh, sn)
+			Ω(err).Should(BeNil())
+			awaitEventPropagation()
+
+			_, err = rmqc.DeleteQueue("/", "mySourceQueue")
 			Ω(err).Should(BeNil())
+			_, err = rmqc.DeleteQueue("/", "myDestQueue")
+			Ω(err).Should(BeNil())
+
+			x, err = rmqc.GetShovel(vh, sn)
+			Ω(x).Should(BeNil())
+			Ω(err).Should(Equal(ErrorResponse{404, "Object Not Found", "Not Found"}))
+		})
+		It("declares a shovel with a numeric delete-after value", func() {
+			vh := "rabbit/hole"
+			sn := "temporary"
+
+			ssu := "amqp://127.0.0.1/%2f"
+			sdu := "amqp://127.0.0.1/%2f"
+
+			shovelDefinition := ShovelDefinition{
+				SourceURI:         ssu,
+				SourceQueue:       "mySourceQueue",
+				DestinationURI:    sdu,
+				DestinationQueue:  "myDestQueue",
+				AddForwardHeaders: true,
+				AckMode:           "on-confirm",
+				DeleteAfter:       "42"}
+
+			_, err := rmqc.DeclareShovel(vh, sn, shovelDefinition)
+			Ω(err).Should(BeNil(), "Error declaring shovel")
+
+			awaitEventPropagation()
+			x, err := rmqc.GetShovel(vh, sn)
+			Ω(err).Should(BeNil(), "Error getting shovel")
 			Ω(x.Name).Should(Equal(sn))
 			Ω(x.Vhost).Should(Equal(vh))
 			Ω(x.Component).Should(Equal("shovel"))
@@ -2436,7 +2482,7 @@ var _ = Describe("Rabbithole", func() {
 			Ω(x.Definition.DestinationQueue).Should(Equal("myDestQueue"))
 			Ω(x.Definition.AddForwardHeaders).Should(Equal(true))
 			Ω(x.Definition.AckMode).Should(Equal("on-confirm"))
-			Ω(x.Definition.DeleteAfter).Should(Equal("never"))
+			Ω(string(x.Definition.DeleteAfter)).Should(Equal("42"))
 
 			_, err = rmqc.DeleteShovel(vh, sn)
 			Ω(err).Should(BeNil())

diff --git a/shovels.go b/shovels.go
@@ -4,6 +4,7 @@ import (
 	"encoding/json"
 	"net/http"
 	"net/url"
+	"strconv"
 )
 
 // ShovelInfo contains the configuration of a shovel
@@ -18,32 +19,65 @@ type ShovelInfo struct {
 	Definition ShovelDefinition `json:"value"`
 }
 
+// DeleteAfter after can hold a delete-after value which may be a string (eg. "never") or an integer
+type DeleteAfter string
+
+// MarshalJSON can marshal a string or an integer
+func (d DeleteAfter) MarshalJSON() ([]byte, error) {
+	deleteAfterInt, err := strconv.Atoi(string(d))
+	if err != nil {
+		return json.Marshal(string(d))
+	}
+	return json.Marshal(deleteAfterInt)
+}
+
+// UnmarshalJSON can unmarshal a string or an integer
+func (d *DeleteAfter) UnmarshalJSON(b []byte) error {
+	// delete-after is a string, such as "never"
+	if b[0] == '"' {
+		var s string
+		if err := json.Unmarshal(b, &s); err != nil {
+			return err
+		}
+		*d = DeleteAfter(s)
+		return nil
+	}
+
+	// delete-after is a number
+	var i int
+	if err := json.Unmarshal(b, &i); err != nil {
+		return err
+	}
+	*d = DeleteAfter(strconv.Itoa(i))
+	return nil
+}
+
 // ShovelDefinition contains the details of the shovel configuration
 type ShovelDefinition struct {
-	AckMode                          string `json:"ack-mode,omitempty"`
-	AddForwardHeaders                bool   `json:"add-forward-headers,omitempty"`
-	DeleteAfter                      string `json:"delete-after,omitempty"`
-	DestinationAddForwardHeaders     bool   `json:"dest-add-forward-headers,omitempty"`
-	DestinationAddTimestampHeader    bool   `json:"dest-add-timestamp-header,omitempty"`
-	DestinationAddress               string `json:"dest-address,omitempty"`
-	DestinationApplicationProperties string `json:"dest-application-properties,omitempty"`
-	DestinationExchange              string `json:"dest-exchange,omitempty"`
-	DestinationExchangeKey           string `json:"dest-exchange-key,omitempty"`
-	DestinationProperties            string `json:"dest-properties,omitempty"`
-	DestinationProtocol              string `json:"dest-protocol,omitempty"`
-	DestinationPublishProperties     string `json:"dest-publish-properties,omitempty"`
-	DestinationQueue                 string `json:"dest-queue,omitempty"`
-	DestinationURI                   string `json:"dest-uri"`
-	PrefetchCount                    int    `json:"prefetch-count,omitempty"`
-	ReconnectDelay                   int    `json:"reconnect-delay,omitempty"`
-	SourceAddress                    string `json:"src-address,omitempty"`
-	SourceDeleteAfter                string `json:"src-delete-after,omitempty"`
-	SourceExchange                   string `json:"src-exchange,omitempty"`
-	SourceExchangeKey                string `json:"src-exchange-key,omitempty"`
-	SourcePrefetchCount              int    `json:"src-prefetch-count,omitempty"`
-	SourceProtocol                   string `json:"src-protocol,omitempty"`
-	SourceQueue                      string `json:"src-queue,omitempty"`
-	SourceURI                        string `json:"src-uri"`
+	AckMode                          string      `json:"ack-mode,omitempty"`
+	AddForwardHeaders                bool        `json:"add-forward-headers,omitempty"`
+	DeleteAfter                      DeleteAfter `json:"delete-after,omitempty"`
+	DestinationAddForwardHeaders     bool        `json:"dest-add-forward-headers,omitempty"`
+	DestinationAddTimestampHeader    bool        `json:"dest-add-timestamp-header,omitempty"`
+	DestinationAddress               string      `json:"dest-address,omitempty"`
+	DestinationApplicationProperties string      `json:"dest-application-properties,omitempty"`
+	DestinationExchange              string      `json:"dest-exchange,omitempty"`
+	DestinationExchangeKey           string      `json:"dest-exchange-key,omitempty"`
+	DestinationProperties            string      `json:"dest-properties,omitempty"`
+	DestinationProtocol              string      `json:"dest-protocol,omitempty"`
+	DestinationPublishProperties     string      `json:"dest-publish-properties,omitempty"`
+	DestinationQueue                 string      `json:"dest-queue,omitempty"`
+	DestinationURI                   string      `json:"dest-uri"`
+	PrefetchCount                    int         `json:"prefetch-count,omitempty"`
+	ReconnectDelay                   int         `json:"reconnect-delay,omitempty"`
+	SourceAddress                    string      `json:"src-address,omitempty"`
+	SourceDeleteAfter                string      `json:"src-delete-after,omitempty"`
+	SourceExchange                   string      `json:"src-exchange,omitempty"`
+	SourceExchangeKey                string      `json:"src-exchange-key,omitempty"`
+	SourcePrefetchCount              int         `json:"src-prefetch-count,omitempty"`
+	SourceProtocol                   string      `json:"src-protocol,omitempty"`
+	SourceQueue                      string      `json:"src-queue,omitempty"`
+	SourceURI                        string      `json:"src-uri"`
 }
 
 // ShovelDefinitionDTO provides a data transfer object