Merge pull request #14 from thejoeejoee/feat-sub-http-response-status

[watermill-http] Custom HTTP response status

pkg/http/context.go

@@ -0,0 +1,30 @@
+package http
+
+import (
+	"context"
+	"github.com/ThreeDotsLabs/watermill/message"
+)
+
+// ctxResponseStatusCodeKey is a context key for the http status code in the message context
+type ctxResponseStatusCodeKey struct{}
+
+// StatusCodeFromContext returns the status code from the context.
+func StatusCodeFromContext(ctx context.Context, otherwise int) int {
+	if v := ctx.Value(ctxResponseStatusCodeKey{}); v != nil {
+		if code, ok := v.(int); ok {
+			return code
+		}
+	}
+	return otherwise
+}
+
+// WithResponseStatusCode returns a new context with the status code.
+func WithResponseStatusCode(ctx context.Context, code int) context.Context {
+	return context.WithValue(ctx, ctxResponseStatusCodeKey{}, code)
+}
+
+// SetResponseStatusCode sets a http status code to the given message.
+func SetResponseStatusCode(m *message.Message, code int) *message.Message {
+	m.SetContext(WithResponseStatusCode(m.Context(), code))
+	return m
+}

pkg/http/pubsub_test.go

@@ -95,6 +95,44 @@ func TestHttpPubSub(t *testing.T) {
 	})
 }
 
+func TestHttpSubStatusCode(t *testing.T) {
+	pub, sub := createPubSub(t)
+
+	defer func() {
+		require.NoError(t, pub.Close())
+		require.NoError(t, sub.Close())
+	}()
+
+	msgs, err := sub.Subscribe(context.Background(), "/test")
+	require.NoError(t, err)
+
+	go func() {
+		_ = sub.StartHTTPServer()
+	}()
+
+	waitForHTTP(t, sub, time.Second*10)
+
+	t.Run("response with custom http status code", func(t *testing.T) {
+		go func() {
+			select {
+			case <-time.After(time.Second * 10):
+				return
+			case msg := <-msgs:
+				http.SetResponseStatusCode(msg, nethttp.StatusForbidden)
+				msg.Nack()
+			}
+		}()
+
+		req, err := nethttp.NewRequest(nethttp.MethodPost, fmt.Sprintf("http://%s/test", sub.Addr()), nil)
+		require.NoError(t, err)
+
+		resp, err := nethttp.DefaultClient.Do(req)
+		require.NoError(t, err)
+
+		require.Equal(t, nethttp.StatusForbidden, resp.StatusCode)
+	})
+}
+
 func waitForHTTP(t *testing.T, sub *http.Subscriber, timeoutTime time.Duration) {
 	timeout := time.After(timeoutTime)
 	for {

pkg/http/subscriber.go

@@ -139,14 +139,17 @@ func (s *Subscriber) Subscribe(ctx context.Context, url string) (<-chan *message
 		s.logger.Trace("Waiting for ACK", logFields)
 		select {
 		case <-msg.Acked():
-			s.logger.Trace("Message acknowledged", logFields.Add(watermill.LogFields{"err": err}))
-			w.WriteHeader(http.StatusOK)
+			code := StatusCodeFromContext(msg.Context(), http.StatusOK)
+			s.logger.Trace("Message acknowledged", logFields.Add(watermill.LogFields{"err": err, "http_status_code": code}))
+			w.WriteHeader(code)
 		case <-msg.Nacked():
-			s.logger.Trace("Message nacked", logFields.Add(watermill.LogFields{"err": err}))
-			w.WriteHeader(http.StatusInternalServerError)
+			code := StatusCodeFromContext(msg.Context(), http.StatusInternalServerError)
+			s.logger.Trace("Message nacked", logFields.Add(watermill.LogFields{"err": err, "http_status_code": code}))
+			w.WriteHeader(code)
 		case <-r.Context().Done():
-			s.logger.Info("Request stopped without ACK received", logFields)
-			w.WriteHeader(http.StatusInternalServerError)
+			code := StatusCodeFromContext(msg.Context(), http.StatusInternalServerError)
+			s.logger.Info("Request stopped without ACK received", logFields.Add(watermill.LogFields{"http_status_code": code}))
+			w.WriteHeader(code)
 		}
 	})