Merge pull request awslabs#96 from youpy/header-set

Use `Header.Set` instead of `Header.Add`

Author: sapessi

core/request.go

@@ -213,13 +213,13 @@ func addToHeader(req *http.Request, apiGwRequest events.APIGatewayProxyRequest)
 		log.Println("Could not marshal stage variables for custom header")
 		return nil, err
 	}
-	req.Header.Add(APIGwStageVarsHeader, string(stageVars))
+	req.Header.Set(APIGwStageVarsHeader, string(stageVars))
 	apiGwContext, err := json.Marshal(apiGwRequest.RequestContext)
 	if err != nil {
 		log.Println("Could not Marshal API GW context for custom header")
 		return req, err
 	}
-	req.Header.Add(APIGwContextHeader, string(apiGwContext))
+	req.Header.Set(APIGwContextHeader, string(apiGwContext))
 	return req, nil
 }
 

core/request_test.go

@@ -217,6 +217,16 @@ var _ = Describe("RequestAccessor tests", func() {
 			// should fail because using header proxy method
 			Expect(ok).To(BeFalse())
 
+			// overwrite existing context header
+			contextRequestWithHeaders := getProxyRequest("orders", "GET")
+			contextRequestWithHeaders.RequestContext = getRequestContext()
+			contextRequestWithHeaders.Headers = map[string]string{core.APIGwContextHeader: `{"AccountID":"abc123"}`}
+			httpReq, err = accessor.ProxyEventToHTTPRequest(contextRequestWithHeaders)
+			Expect(err).To(BeNil())
+			headerContext, err = accessor.GetAPIGatewayContext(httpReq)
+			Expect(err).To(BeNil())
+			Expect(headerContext.AccountID).To(Equal("x"))
+
 			httpReq, err = accessor.EventToRequestWithContext(context.Background(), contextRequest)
 			Expect(err).To(BeNil())
 			proxyContext, ok = core.GetAPIGatewayContextFromContext(httpReq.Context())
@@ -264,6 +274,16 @@ var _ = Describe("RequestAccessor tests", func() {
 			Expect("value1").To(Equal(stageVars["var1"]))
 			Expect("value2").To(Equal(stageVars["var2"]))
 
+			// overwrite existing stagevars header
+			varsRequestWithHeaders := getProxyRequest("orders", "GET")
+			varsRequestWithHeaders.StageVariables = getStageVariables()
+			varsRequestWithHeaders.Headers = map[string]string{core.APIGwStageVarsHeader: `{"var1":"abc123"}`}
+			httpReq, err = accessor.ProxyEventToHTTPRequest(varsRequestWithHeaders)
+			Expect(err).To(BeNil())
+			stageVars, err = accessor.GetAPIGatewayStageVars(httpReq)
+			Expect(err).To(BeNil())
+			Expect(stageVars["var1"]).To(Equal("value1"))
+
 			stageVars, ok := core.GetStageVarsFromContext(httpReq.Context())
 			// not present in context
 			Expect(ok).To(BeFalse())