simplify and fix the json middleware example

Author: robbyt

examples/README.md

@@ -17,6 +17,5 @@ HTTP server cluster with automatic port allocation.
 ## Running
 
 ```bash
-cd examples/<name>
-go run main.go
+go run ./examples/<name>
 ```

examples/custom_middleware/example/jsonenforcer.go

@@ -8,66 +8,64 @@ import (
 	"github.com/robbyt/go-supervisor/runnables/httpserver"
 )
 
-// ResponseWrapper captures response data for transformation
+// ResponseBuffer captures response data for transformation
 //
 // This is a simple example for demonstration purposes and is not intended for
 // production use. Limitations:
 // - Does not preserve optional HTTP interfaces (http.Hijacker, http.Flusher, http.Pusher)
 // - Not safe for concurrent writes from multiple goroutines within the same request
 // - No memory limits on buffered content
 //
-// Each request gets its own ResponseWrapper instance, so different requests won't
+// Each request gets its own ResponseBuffer instance, so different requests won't
 // interfere with each other.
-type ResponseWrapper struct {
-	original httpserver.ResponseWriter
-	buffer   *bytes.Buffer
-	headers  http.Header
-	status   int
+type ResponseBuffer struct {
+	buffer  *bytes.Buffer
+	headers http.Header
+	status  int
 }
 
-// NewResponseWrapper creates a new response wrapper
-func NewResponseWrapper(original httpserver.ResponseWriter) *ResponseWrapper {
-	return &ResponseWrapper{
-		original: original,
-		buffer:   new(bytes.Buffer),
-		headers:  make(http.Header),
-		status:   0, // 0 means not set yet
+// NewResponseBuffer creates a new response buffer
+func NewResponseBuffer() *ResponseBuffer {
+	return &ResponseBuffer{
+		buffer:  new(bytes.Buffer),
+		headers: make(http.Header),
+		status:  0, // 0 means not set yet
 	}
 }
 
 // Header implements http.ResponseWriter
-func (rw *ResponseWrapper) Header() http.Header {
-	return rw.headers
+func (rb *ResponseBuffer) Header() http.Header {
+	return rb.headers
 }
 
 // Write implements http.ResponseWriter
-func (rw *ResponseWrapper) Write(data []byte) (int, error) {
-	return rw.buffer.Write(data)
+func (rb *ResponseBuffer) Write(data []byte) (int, error) {
+	return rb.buffer.Write(data)
 }
 
 // WriteHeader implements http.ResponseWriter
-func (rw *ResponseWrapper) WriteHeader(statusCode int) {
-	if rw.status == 0 {
-		rw.status = statusCode
+func (rb *ResponseBuffer) WriteHeader(statusCode int) {
+	if rb.status == 0 {
+		rb.status = statusCode
 	}
 }
 
 // Status implements httpserver.ResponseWriter
-func (rw *ResponseWrapper) Status() int {
-	if rw.status == 0 && rw.buffer.Len() > 0 {
+func (rb *ResponseBuffer) Status() int {
+	if rb.status == 0 && rb.buffer.Len() > 0 {
 		return http.StatusOK
 	}
-	return rw.status
+	return rb.status
 }
 
 // Written implements httpserver.ResponseWriter
-func (rw *ResponseWrapper) Written() bool {
-	return rw.buffer.Len() > 0 || rw.status != 0
+func (rb *ResponseBuffer) Written() bool {
+	return rb.buffer.Len() > 0 || rb.status != 0
 }
 
 // Size implements httpserver.ResponseWriter
-func (rw *ResponseWrapper) Size() int {
-	return rw.buffer.Len()
+func (rb *ResponseBuffer) Size() int {
+	return rb.buffer.Len()
 }
 
 // transformToJSON wraps non-JSON content in a JSON response
@@ -90,23 +88,25 @@ func transformToJSON(data []byte) ([]byte, error) {
 // Valid JSON responses are preserved as-is.
 func New() httpserver.HandlerFunc {
 	return func(rp *httpserver.RequestProcessor) {
-		// Wrap the response writer to capture output
-		wrapper := NewResponseWrapper(rp.Writer())
-		rp.SetWriter(wrapper)
+		// Store original writer before buffering
+		originalWriter := rp.Writer()
+
+		// Buffer the response to capture output
+		buffer := NewResponseBuffer()
+		rp.SetWriter(buffer)
 
 		// Continue to next middleware/handler
 		rp.Next()
 
 		// RESPONSE PHASE: Transform response to JSON
-		originalData := wrapper.buffer.Bytes()
-		statusCode := wrapper.Status()
+		originalData := buffer.buffer.Bytes()
+		statusCode := buffer.Status()
 		if statusCode == 0 {
 			statusCode = http.StatusOK
 		}
 
 		// Copy headers to original writer
-		originalWriter := wrapper.original
-		for key, values := range wrapper.Header() {
+		for key, values := range buffer.Header() {
 			for _, value := range values {
 				originalWriter.Header().Add(key, value)
 			}
@@ -120,7 +120,7 @@ func New() httpserver.HandlerFunc {
 		}
 
 		// Transform captured data to JSON
-		if len(originalData) == 0 && wrapper.status == 0 {
+		if len(originalData) == 0 && buffer.status == 0 {
 			return
 		}
 

examples/custom_middleware/example/jsonenforcer_test.go

@@ -9,6 +9,18 @@ import (
 	"github.com/stretchr/testify/assert"
 )
 
+func createTestRoute(t *testing.T, handler http.HandlerFunc) *httpserver.Route {
+	t.Helper()
+	route, err := httpserver.NewRouteFromHandlerFunc(
+		"test",
+		"/test",
+		handler,
+		New(),
+	)
+	assert.NoError(t, err)
+	return route
+}
+
 func TestJSONEnforcerMiddleware(t *testing.T) {
 	t.Parallel()
 
@@ -22,14 +34,7 @@ func TestJSONEnforcerMiddleware(t *testing.T) {
 			assert.NoError(t, err)
 		}
 
-		// Create route with JSON enforcer middleware
-		route, err := httpserver.NewRouteFromHandlerFunc(
-			"test",
-			"/test",
-			textHandler,
-			New(),
-		)
-		assert.NoError(t, err)
+		route := createTestRoute(t, textHandler)
 
 		route.ServeHTTP(rec, req)
 
@@ -53,13 +58,7 @@ func TestJSONEnforcerMiddleware(t *testing.T) {
 			assert.NoError(t, err)
 		}
 
-		route, err := httpserver.NewRouteFromHandlerFunc(
-			"test",
-			"/test",
-			jsonHandler,
-			New(),
-		)
-		assert.NoError(t, err)
+		route := createTestRoute(t, jsonHandler)
 
 		route.ServeHTTP(rec, req)
 
@@ -83,13 +82,7 @@ func TestJSONEnforcerMiddleware(t *testing.T) {
 			assert.NoError(t, err)
 		}
 
-		route, err := httpserver.NewRouteFromHandlerFunc(
-			"test",
-			"/test",
-			errorHandler,
-			New(),
-		)
-		assert.NoError(t, err)
+		route := createTestRoute(t, errorHandler)
 
 		route.ServeHTTP(rec, req)
 
@@ -112,13 +105,7 @@ func TestJSONEnforcerMiddleware(t *testing.T) {
 			w.WriteHeader(http.StatusNoContent)
 		}
 
-		route, err := httpserver.NewRouteFromHandlerFunc(
-			"test",
-			"/test",
-			emptyHandler,
-			New(),
-		)
-		assert.NoError(t, err)
+		route := createTestRoute(t, emptyHandler)
 
 		route.ServeHTTP(rec, req)
 
@@ -139,13 +126,7 @@ func TestJSONEnforcerMiddleware(t *testing.T) {
 			assert.NoError(t, err)
 		}
 
-		route, err := httpserver.NewRouteFromHandlerFunc(
-			"test",
-			"/test",
-			notModifiedHandler,
-			New(),
-		)
-		assert.NoError(t, err)
+		route := createTestRoute(t, notModifiedHandler)
 
 		route.ServeHTTP(rec, req)
 
@@ -164,13 +145,7 @@ func TestJSONEnforcerMiddleware(t *testing.T) {
 			assert.NoError(t, err)
 		}
 
-		route, err := httpserver.NewRouteFromHandlerFunc(
-			"test",
-			"/test",
-			htmlHandler,
-			New(),
-		)
-		assert.NoError(t, err)
+		route := createTestRoute(t, htmlHandler)
 
 		route.ServeHTTP(rec, req)
 
@@ -192,13 +167,7 @@ func TestJSONEnforcerMiddleware(t *testing.T) {
 			assert.NoError(t, err)
 		}
 
-		route, err := httpserver.NewRouteFromHandlerFunc(
-			"test",
-			"/test",
-			arrayHandler,
-			New(),
-		)
-		assert.NoError(t, err)
+		route := createTestRoute(t, arrayHandler)
 
 		route.ServeHTTP(rec, req)
 
@@ -222,13 +191,7 @@ func TestJSONEnforcerMiddleware(t *testing.T) {
 			assert.NoError(t, err)
 		}
 
-		route, err := httpserver.NewRouteFromHandlerFunc(
-			"test",
-			"/test",
-			malformedHandler,
-			New(),
-		)
-		assert.NoError(t, err)
+		route := createTestRoute(t, malformedHandler)
 
 		route.ServeHTTP(rec, req)
 
@@ -255,13 +218,7 @@ func TestJSONEnforcerMiddleware(t *testing.T) {
 			assert.NoError(t, err)
 		}
 
-		route, err := httpserver.NewRouteFromHandlerFunc(
-			"test",
-			"/test",
-			multiWriteHandler,
-			New(),
-		)
-		assert.NoError(t, err)
+		route := createTestRoute(t, multiWriteHandler)
 
 		route.ServeHTTP(rec, req)
 
@@ -274,33 +231,39 @@ func TestJSONEnforcerMiddleware(t *testing.T) {
 			"all writes should be buffered and wrapped",
 		)
 	})
-}
-
-// Test compliance with the middleware framework
-func TestJSONEnforcerCompliance(t *testing.T) {
-	// This is a placeholder test that would import the compliance framework
-	// In practice, you would add this to runnables/httpserver/middleware/compliance_test.go
-	jsonMiddleware := New()
 
-	// Test basic compliance: calls Next()
-	t.Run("calls Next()", func(t *testing.T) {
+	t.Run("preserves headers from handler", func(t *testing.T) {
 		req := httptest.NewRequest("GET", "/test", nil)
 		rec := httptest.NewRecorder()
 
-		called := false
-		mockHandler := func(w http.ResponseWriter, r *http.Request) {
-			called = true
+		headerHandler := func(w http.ResponseWriter, r *http.Request) {
+			w.Header().Set("X-Custom-Header", "test-value")
+			w.Header().Set("X-Another-Header", "another-value")
+			_, err := w.Write([]byte("Hello World"))
+			assert.NoError(t, err)
 		}
 
-		route, err := httpserver.NewRouteFromHandlerFunc(
-			"test",
-			"/test",
-			mockHandler,
-			jsonMiddleware,
-		)
-		assert.NoError(t, err)
+		route := createTestRoute(t, headerHandler)
+		route.ServeHTTP(rec, req)
 
+		assert.Equal(t, "application/json", rec.Header().Get("Content-Type"))
+		assert.Equal(t, "test-value", rec.Header().Get("X-Custom-Header"))
+		assert.Equal(t, "another-value", rec.Header().Get("X-Another-Header"))
+		assert.Contains(t, rec.Body.String(), `{"response":"Hello World"}`)
+	})
+
+	t.Run("handles empty response", func(t *testing.T) {
+		req := httptest.NewRequest("GET", "/test", nil)
+		rec := httptest.NewRecorder()
+
+		emptyHandler := func(w http.ResponseWriter, r *http.Request) {
+			// Handler does nothing - no writes, no status
+		}
+
+		route := createTestRoute(t, emptyHandler)
 		route.ServeHTTP(rec, req)
-		assert.True(t, called, "should call the next handler")
+
+		assert.Equal(t, http.StatusOK, rec.Code)
+		assert.Empty(t, rec.Body.String(), "empty response should remain empty")
 	})
 }