Retry middleware : store headers per attempts and propagate them when responding.

old/middlewares/retry_test.go

@@ -1,8 +1,10 @@
 package middlewares
 
 import (
+	"fmt"
 	"net/http"
 	"net/http/httptest"
+	"net/http/httptrace"
 	"strings"
 	"testing"
 
@@ -258,3 +260,45 @@ func TestRetryWithFlush(t *testing.T) {
 		t.Errorf("Wrong body %q want %q", responseRecorder.Body.String(), "FULL DATA")
 	}
 }
+
+func TestMultipleRetriesShouldNotLooseHeaders(t *testing.T) {
+	attempt := 0
+	expectedHeaderName := "X-Foo-Test-2"
+	expectedHeaderValue := "bar"
+
+	next := http.HandlerFunc(func(rw http.ResponseWriter, req *http.Request) {
+		headerName := fmt.Sprintf("X-Foo-Test-%d", attempt)
+		rw.Header().Add(headerName, expectedHeaderValue)
+		if attempt < 2 {
+			attempt++
+			return
+		}
+
+		// Request has been successfully written to backend
+		trace := httptrace.ContextClientTrace(req.Context())
+		trace.WroteHeaders()
+
+		// And we decide to answer to client
+		rw.WriteHeader(http.StatusNoContent)
+	})
+
+	retry := NewRetry(3, next, &countingRetryListener{})
+	responseRecorder := httptest.NewRecorder()
+	retry.ServeHTTP(responseRecorder, &http.Request{})
+
+	headerValue := responseRecorder.Header().Get(expectedHeaderName)
+
+	// Validate if we have the correct header
+	if headerValue != expectedHeaderValue {
+		t.Errorf("Expected to have %s for header %s, got %s", expectedHeaderValue, expectedHeaderName, headerValue)
+	}
+
+	// Validate that we don't have headers from previous attempts
+	for i := 0; i < attempt; i++ {
+		headerName := fmt.Sprintf("X-Foo-Test-%d", i)
+		headerValue = responseRecorder.Header().Get("headerName")
+		if headerValue != "" {
+			t.Errorf("Expected no value for header %s, got %s", headerName, headerValue)
+		}
+	}
+}