Apply comments

Author: ataleksandrov

api_test.go

@@ -26,6 +26,8 @@ import (
 	"net/url"
 	"strings"
 
+	"github.com/onsi/ginkgo/extensions/table"
+
 	"github.com/pivotal-cf/brokerapi/middlewares"
 
 	"code.cloudfoundry.org/lager"
@@ -111,13 +113,6 @@ var _ = Describe("Service Broker API", func() {
 			return recorder
 		}
 
-		It("has a X-Correlation-ID header", func() {
-			response := makeRequest()
-
-			header := response.Header().Get("X-Correlation-ID")
-			Expect(header).Should(Not(BeNil()))
-		})
-
 		It("has a Content-Type header", func() {
 			response := makeRequest()
 
@@ -2337,6 +2332,11 @@ var _ = Describe("Service Broker API", func() {
 	})
 
 	Describe("CorrelationIDHeader", func() {
+		const correlationID = "fake-correlation-id"
+
+		type testCase struct {
+			correlationIDHeaderName string
+		}
 
 		var (
 			fakeServiceBroker *fakes.AutoFakeServiceBroker
@@ -2360,20 +2360,33 @@ var _ = Describe("Service Broker API", func() {
 			testServer.Close()
 		})
 
-		When("X-Correlation-ID is passed", func() {
-			It("Adds correlation id to the context", func() {
-				const correlationID = "fake-correlation-id"
-				req.Header.Add("X-Correlation-ID", correlationID)
+		table.DescribeTable("Adds correlation id to the context", func(tc testCase) {
+			req.Header.Add(tc.correlationIDHeaderName, correlationID)
 
-				_, err := http.DefaultClient.Do(req)
-				Expect(err).NotTo(HaveOccurred())
+			_, err := http.DefaultClient.Do(req)
+			Expect(err).NotTo(HaveOccurred())
 
-				Expect(fakeServiceBroker.ServicesCallCount()).To(Equal(1), "Services was not called")
-				ctx := fakeServiceBroker.ServicesArgsForCall(0)
-				Expect(ctx.Value(middlewares.CorrelationIDKey)).To(Equal(correlationID))
+			Expect(fakeServiceBroker.ServicesCallCount()).To(Equal(1), "Services was not called")
+			ctx := fakeServiceBroker.ServicesArgsForCall(0)
+			Expect(ctx.Value(middlewares.CorrelationIDKey)).To(Equal(correlationID))
+		},
+			table.Entry("X-Correlation-ID", testCase{
+				correlationIDHeaderName: "X-Correlation-ID",
+			}),
+			table.Entry("X-CorrelationID", testCase{
+				correlationIDHeaderName: "X-CorrelationID",
+			}),
+			table.Entry("X-ForRequest-ID", testCase{
+				correlationIDHeaderName: "X-ForRequest-ID",
+			}),
+			table.Entry("X-Request-ID", testCase{
+				correlationIDHeaderName: "X-Request-ID",
+			}),
+			table.Entry("X-Vcap-Request-Id", testCase{
+				correlationIDHeaderName: "X-Vcap-Request-Id",
+			}),
+		)
 
-			})
-		})
 		When("X-Correlation-ID is not passed", func() {
 			It("Generates correlation id and adds it to the context", func() {
 				_, err := http.DefaultClient.Do(req)

middlewares/correlation_id_header.go

@@ -13,25 +13,22 @@ var correlationIDHeaders = []string{"X-Correlation-ID", "X-CorrelationID", "X-Fo
 
 func AddCorrelationIDToContext(next http.Handler) http.Handler {
 	return http.HandlerFunc(func(w http.ResponseWriter, req *http.Request) {
-		var correlationID, headerName string
+		var correlationID string
 		var found bool
 
 		for _, header := range correlationIDHeaders {
 			headerValue := req.Header.Get(header)
 			if headerValue != "" {
 				correlationID = headerValue
-				headerName = header
 				found = true
 				break
 			}
 		}
 
 		if !found {
 			correlationID = generateCorrelationID()
-			headerName = correlationIDHeaders[0]
 		}
 
-		w.Header().Set(headerName, correlationID)
 		newCtx := context.WithValue(req.Context(), CorrelationIDKey, correlationID)
 		next.ServeHTTP(w, req.WithContext(newCtx))
 	})