X-Ray support improved

gluo_test.go

@@ -47,7 +47,8 @@ const testRequest = `{
             "sourceIp":"127.0.0.1",
             "accountId":null
         },
-        "stage":"prod"
+        "stage":"prod",
+        "X-Amzn-Trace-Id":"Root=1-5759e988-bd862e3fe1be46a994272793;Sampled=1"
     },
     "queryStringParameters":{
         "foo":"bar"
@@ -66,7 +67,6 @@ const testRequest = `{
         "Host":"1234567890.execute-api.us-east-1.amazonaws.com",
         "X-Forwarded-Proto":"https",
         "X-Amz-Cf-Id":"cDehVQoZnx43VYQb9j2-nvCh-9z396Uhbp027Y2JvkCPNLmGJHqlaA==",
-        "X-Amz-Trace-Id":"Root=1-5759e988-bd862e3fe1be46a994272793;Sampled=1",
         "CloudFront-Is-Tablet-Viewer":"false",
         "Cache-Control":"max-age=0",
         "User-Agent":"Custom User Agent String",
@@ -90,6 +90,7 @@ type helloRequest struct {
 func TestLambdaServe(t *testing.T) {
 	la := LambdaAdapter{
 		http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+			w.Header().Set("Content-Type", "application/json")
 			buffer := new(bytes.Buffer)
 			buffer.ReadFrom(r.Body)
 			r.Body.Close()
@@ -104,12 +105,13 @@ func TestLambdaServe(t *testing.T) {
 				hr.Name = "stranger"
 			}
 			result := fmt.Sprintf("Hello, %s.", hr.Name)
-			w.Header().Set("Content-Type", "application/json")
 			w.Write([]byte(result))
 		}),
 	}
 	req := events.APIGatewayProxyRequest{}
-	json.Unmarshal([]byte(testRequest), &req)
+	if err := json.Unmarshal([]byte(testRequest), &req); err != nil {
+		t.Errorf("unexpected error on json.Unmarshal: %v", err)
+	}
 	res, err := la.Handle(context.TODO(), req)
 	if err != nil {
 		t.Error("LambdaHandler.Handle must return nil")

request.go

@@ -54,8 +54,8 @@ func request(ctx context.Context, event events.APIGatewayProxyRequest) (*http.Re
 	req.Header.Set("X-Request-ID", event.RequestContext.RequestID)
 	req.Header.Set("X-Stage", event.RequestContext.Stage)
 	// AWS X-Ray
-	if traceID := ctx.Value("X-Amzn-Trace-ID"); traceID != nil {
-		req.Header.Set("X-Amzn-Trace-ID", fmt.Sprintf("%v", traceID))
+	if traceID := ctx.Value("X-Amzn-Trace-Id"); traceID != nil {
+		req.Header.Set("X-Amzn-Trace-Id", fmt.Sprintf("%v", traceID))
 	}
 	req.RemoteAddr = event.RequestContext.Identity.SourceIP
 	return req, nil

request_test.go

@@ -2,7 +2,9 @@ package gluo
 
 import (
 	"bytes"
+	"context"
 	"encoding/base64"
+	"encoding/json"
 	"github.com/aws/aws-lambda-go/events"
 	"testing"
 )
@@ -49,3 +51,16 @@ func TestWrongBase64Body(t *testing.T) {
 		t.Errorf("expected error reading body in base 64 from getBodyReader, got nil")
 	}
 }
+
+func TestXRayHeader(t *testing.T) {
+	event := events.APIGatewayProxyRequest{}
+	json.Unmarshal([]byte(testRequest), &event)
+	ctx := context.WithValue(context.Background(), "X-Amzn-Trace-Id", "Root=1-5759e988-bd862e3fe1be46a994272793;Sampled=1")
+	req, err := request(ctx, event)
+	if err != nil {
+		t.Errorf("unexpected error on request: %v", err)
+	}
+	if req.Header.Get("X-Amzn-Trace-Id") == "" {
+		t.Errorf("expected X-Amzn-Trace-Id injected but it wasn't found")
+	}
+}