middlewares: fix RequestIDHandler to add trackingID on response heade…

…rs (#11)

* extracted "X-Request-ID" header to constants.HeaderKeyRequestID
* add response header X-Request-ID with requestID
* added tests

constants/constants.go

@@ -13,5 +13,7 @@ const (
 	LogKeyResquestID       = "_request_id"
 	LogKeyTimestamp        = "_timestamp"
 	LogKeyURLPath          = "_url_path"
-	LogKeyUserAgent        = "_user_agent"	
+	LogKeyUserAgent        = "_user_agent"
+
+	HeaderKeyRequestID = "X-Request-ID"
 )

middlewares/request_id.go

@@ -4,21 +4,23 @@ import (
 	"context"
 	"net/http"
 
+	"github.com/AndersonQ/go-skeleton/constants"
+
 	"github.com/google/uuid"
 )
 
 type reqIDCtx int
 
-const reqIDContexKey reqIDCtx = reqIDCtx(0)
+const reaIDContextKey = reqIDCtx(0)
 
 // requestIDContext creates a context with request id
 func requestIDContext(ctx context.Context, rid string) context.Context {
-	return context.WithValue(ctx, reqIDContexKey, rid)
+	return context.WithValue(ctx, reaIDContextKey, rid)
 }
 
 // requestIDFromContext returns the request id from context
 func requestIDFromContext(ctx context.Context) string {
-	rid, ok := ctx.Value(reqIDContexKey).(string)
+	rid, ok := ctx.Value(reaIDContextKey).(string)
 	if !ok {
 		return ""
 	}
@@ -30,11 +32,14 @@ func requestIDFromContext(ctx context.Context) string {
 // request id. Otherwise, generates a new unique ID.
 func RequestIDHandler(h http.Handler) http.Handler {
 	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
-		rid := r.Header.Get("X-Request-ID")
+		rid := r.Header.Get(constants.HeaderKeyRequestID)
+
 		if rid == "" {
 			rid = uuid.New().String()
-			r.Header.Set("X-Request-ID", rid)
 		}
+
+		w.Header().Set(constants.HeaderKeyRequestID, rid)
+
 		ctx := requestIDContext(r.Context(), rid)
 		h.ServeHTTP(w, r.WithContext(ctx))
 	})

middlewares/request_id_test.go

@@ -0,0 +1,57 @@
+package middlewares
+
+import (
+	"net/http"
+	"net/http/httptest"
+	"testing"
+
+	"github.com/AndersonQ/go-skeleton/constants"
+)
+
+func TestRequestIDHandlerWithRequestID(t *testing.T) {
+	want := "requestID"
+
+	r := httptest.NewRequest("ignore", "/ignore", nil)
+	r.Header.Set(constants.HeaderKeyRequestID, want)
+	w := httptest.NewRecorder()
+
+	h := RequestIDHandler(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {}))
+
+	h.ServeHTTP(w, r)
+
+	got := w.Header().Get(constants.HeaderKeyRequestID)
+
+	if got != want {
+		t.Errorf("got: %s, want: %s", got, want)
+	}
+}
+
+func TestRequestIDHandlerWithoutRequestID(t *testing.T) {
+	r := httptest.NewRequest("ignore", "/ignore", nil)
+	w := httptest.NewRecorder()
+	var want string
+
+	h := RequestIDHandler(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		got, ok := r.Context().Value(reaIDContextKey).(string)
+
+		if !ok {
+			t.Fatalf("requestID is not a string! got: %s", got)
+		}
+		if got == "" {
+			t.Fatalf("requestID from request context is empty, want a uuid as requestID")
+		}
+
+		want = got
+	}))
+
+	h.ServeHTTP(w, r)
+	got := w.Header().Get(constants.HeaderKeyRequestID)
+
+	if got == "" {
+		t.Errorf("got: %s, want an uuid", got)
+	}
+
+	if got != want {
+		t.Errorf("got requestID %s, want: %s", got, got)
+	}
+}