Merge pull request #97 from remind101/http-handler-compat

Ensure Context arg is the same as request Context
remind101 · Apr 4, 2018 · e2ec3e8 · e2ec3e8
2 parents a2aea2a + 8f46cf0
commit e2ec3e8
Show file tree

Hide file tree

Showing 7 changed files with 23 additions and 1 deletion.
diff --git a/httpx/middleware/header.go b/httpx/middleware/header.go
@@ -28,6 +28,8 @@ func (h *Header) ServeHTTPContext(ctx context.Context, w http.ResponseWriter, r
 	value := e(r)
 
 	ctx = httpx.WithHeader(ctx, h.key, value)
+	r = r.WithContext(ctx)
+
 	return h.handler.ServeHTTPContext(ctx, w, r)
 }
 

diff --git a/httpx/middleware/header_test.go b/httpx/middleware/header_test.go
@@ -6,6 +6,7 @@ import (
 	"testing"
 
 	"context"
+
 	"github.com/remind101/pkg/httpx"
 )
 
@@ -24,7 +25,11 @@ func TestHeader(t *testing.T) {
 		m := ExtractHeader(
 			httpx.HandlerFunc(func(ctx context.Context, w http.ResponseWriter, r *http.Request) error {
 				data := httpx.Header(ctx, tt.key)
+				if got, want := data, tt.val; got != want {
+					t.Fatalf("%s => %s; want %s", tt.key, got, want)
+				}
 
+				data = httpx.Header(r.Context(), tt.key)
 				if got, want := data, tt.val; got != want {
 					t.Fatalf("%s => %s; want %s", tt.key, got, want)
 				}

diff --git a/httpx/middleware/logger.go b/httpx/middleware/logger.go
@@ -43,7 +43,10 @@ func LogTo(h httpx.Handler, g loggerGenerator) httpx.Handler {
 func InsertLogger(h httpx.Handler, g loggerGenerator) httpx.Handler {
 	return httpx.HandlerFunc(func(ctx context.Context, w http.ResponseWriter, r *http.Request) error {
 		l := g(ctx, r)
+
 		ctx = logger.WithLogger(ctx, l)
+		r = r.WithContext(ctx)
+
 		return h.ServeHTTPContext(ctx, w, r)
 	})
 }

diff --git a/httpx/middleware/opentracing.go b/httpx/middleware/opentracing.go
@@ -45,6 +45,7 @@ func (h *OpentracingTracer) ServeHTTPContext(ctx context.Context, w http.Respons
 
 	defer span.Finish()
 	ctx = opentracing.ContextWithSpan(ctx, span)
+	r = r.WithContext(ctx)
 
 	rw := NewResponseWriter(w)
 	reqErr := h.handler.ServeHTTPContext(ctx, rw, r)

diff --git a/httpx/middleware/reporter.go b/httpx/middleware/reporter.go
@@ -26,6 +26,8 @@ func (m *Reporter) ServeHTTPContext(ctx context.Context, w http.ResponseWriter,
 	// Add the request id to reporter context.
 	ctx = errors.WithInfo(ctx, "request_id", httpx.RequestID(ctx))
 
+	r = r.WithContext(ctx)
+
 	return m.handler.ServeHTTPContext(ctx, w, r)
 }
 

diff --git a/httpx/middleware/request_id.go b/httpx/middleware/request_id.go
@@ -4,6 +4,7 @@ import (
 	"net/http"
 
 	"context"
+
 	"github.com/remind101/pkg/httpx"
 )
 
@@ -39,5 +40,7 @@ func (h *RequestID) ServeHTTPContext(ctx context.Context, w http.ResponseWriter,
 	requestID := e(r)
 
 	ctx = httpx.WithRequestID(ctx, requestID)
+	r = r.WithContext(ctx)
+
 	return h.handler.ServeHTTPContext(ctx, w, r)
 }
diff --git a/httpx/middleware/request_id_test.go b/httpx/middleware/request_id_test.go
@@ -5,8 +5,9 @@ import (
 	"net/http/httptest"
 	"testing"
 
-	"github.com/remind101/pkg/httpx"
 	"context"
+
+	"github.com/remind101/pkg/httpx"
 )
 
 func TestRequestID(t *testing.T) {
@@ -23,7 +24,12 @@ func TestRequestID(t *testing.T) {
 		m := &RequestID{
 			handler: httpx.HandlerFunc(func(ctx context.Context, w http.ResponseWriter, r *http.Request) error {
 				requestID := httpx.RequestID(ctx)
+				if got, want := requestID, tt.id; got != want {
+					t.Fatalf("RequestID => %s; want %s", got, want)
+				}
 
+				// From request.Context()
+				requestID = httpx.RequestID(r.Context())
 				if got, want := requestID, tt.id; got != want {
 					t.Fatalf("RequestID => %s; want %s", got, want)
 				}