Test uses trace context from context object.

Author: purple4reina

internal/trace/context_test.go

@@ -14,10 +14,9 @@ import (
 	"io/ioutil"
 	"testing"
 
+	"github.com/DataDog/datadog-lambda-go/internal/extension"
 	"github.com/aws/aws-xray-sdk-go/header"
-
 	"github.com/aws/aws-xray-sdk-go/xray"
-
 	"github.com/stretchr/testify/assert"
 )
 
@@ -45,6 +44,20 @@ func mockLambdaXRayTraceContext(ctx context.Context, traceID, parentID string, s
 	return context.WithValue(ctx, xray.LambdaTraceHeaderKey, headerString)
 }
 
+func mockTraceContext(traceID, parentID, samplingPriority string) context.Context {
+	ctx := context.Background()
+	if traceID != "" {
+		ctx = context.WithValue(ctx, extension.DdTraceId, traceID)
+	}
+	if parentID != "" {
+		ctx = context.WithValue(ctx, extension.DdParentId, parentID)
+	}
+	if samplingPriority != "" {
+		ctx = context.WithValue(ctx, extension.DdSamplingPriority, samplingPriority)
+	}
+	return ctx
+}
+
 func loadRawJSON(t *testing.T, filename string) *json.RawMessage {
 	bytes, err := ioutil.ReadFile(filename)
 	if err != nil {
@@ -117,6 +130,63 @@ func TestGetDatadogTraceContextForMissingData(t *testing.T) {
 	assert.False(t, ok)
 }
 
+func TestGetDatadogTraceContextFromContextObject(t *testing.T) {
+	testcases := []struct {
+		traceID          string
+		parentID         string
+		samplingPriority string
+		expectTC         TraceContext
+		expectOk         bool
+	}{
+		{
+			"trace",
+			"parent",
+			"sampling",
+			TraceContext{
+				"x-datadog-trace-id":          "trace",
+				"x-datadog-parent-id":         "parent",
+				"x-datadog-sampling-priority": "sampling",
+			},
+			true,
+		},
+		{
+			"",
+			"parent",
+			"sampling",
+			TraceContext{},
+			false,
+		},
+		{
+			"trace",
+			"",
+			"sampling",
+			TraceContext{},
+			false,
+		},
+		{
+			"trace",
+			"parent",
+			"",
+			TraceContext{
+				"x-datadog-trace-id":          "trace",
+				"x-datadog-parent-id":         "parent",
+				"x-datadog-sampling-priority": "1",
+			},
+			true,
+		},
+	}
+
+	ev := loadRawJSON(t, "../testdata/non-proxy-no-headers.json")
+	for _, test := range testcases {
+		t.Run(test.traceID+test.parentID+test.samplingPriority, func(t *testing.T) {
+			ctx := mockTraceContext(test.traceID, test.parentID, test.samplingPriority)
+			tc, ok := getTraceContext(ctx, getHeadersFromEventHeaders(ctx, *ev))
+			assert.Equal(t, test.expectTC, tc)
+			assert.Equal(t, test.expectOk, ok)
+		})
+	}
+}
+
 func TestConvertXRayTraceID(t *testing.T) {
 	output, err := convertXRayTraceIDToDatadogTraceID(mockXRayTraceID)
 	assert.NoError(t, err)