diff --git a/contrib/aws/aws-sdk-go-v2/aws/aws_test.go b/contrib/aws/aws-sdk-go-v2/aws/aws_test.go index 88fba42c5d..6a5412a175 100644 --- a/contrib/aws/aws-sdk-go-v2/aws/aws_test.go +++ b/contrib/aws/aws-sdk-go-v2/aws/aws_test.go @@ -8,6 +8,7 @@ package aws import ( "context" "encoding/base64" + "encoding/json" "net/http" "net/http/httptest" "net/url" @@ -24,12 +25,13 @@ import ( "github.com/aws/aws-sdk-go-v2/service/dynamodb" "github.com/aws/aws-sdk-go-v2/service/ec2" "github.com/aws/aws-sdk-go-v2/service/eventbridge" + eventBridgeTypes "github.com/aws/aws-sdk-go-v2/service/eventbridge/types" "github.com/aws/aws-sdk-go-v2/service/kinesis" "github.com/aws/aws-sdk-go-v2/service/s3" "github.com/aws/aws-sdk-go-v2/service/sfn" "github.com/aws/aws-sdk-go-v2/service/sns" "github.com/aws/aws-sdk-go-v2/service/sqs" - "github.com/aws/aws-sdk-go-v2/service/sqs/types" + sqsTypes "github.com/aws/aws-sdk-go-v2/service/sqs/types" "github.com/aws/smithy-go/middleware" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" @@ -281,6 +283,66 @@ func TestAppendMiddlewareSqsReceiveMessage(t *testing.T) { } } +func TestAppendMiddlewareSqsSendMessage(t *testing.T) { + mt := mocktracer.Start() + defer mt.Stop() + + expectedStatusCode := 200 + server := mockAWS(expectedStatusCode) + defer server.Close() + + resolver := aws.EndpointResolverFunc(func(service, region string) (aws.Endpoint, error) { + return aws.Endpoint{ + PartitionID: "aws", + URL: server.URL, + SigningRegion: "eu-west-1", + }, nil + }) + + awsCfg := aws.Config{ + Region: "eu-west-1", + Credentials: aws.AnonymousCredentials{}, + EndpointResolver: resolver, + } + + AppendMiddleware(&awsCfg) + + sqsClient := sqs.NewFromConfig(awsCfg) + sendMessageInput := &sqs.SendMessageInput{ + MessageBody: aws.String("test message"), + QueueUrl: aws.String("https://sqs.us-west-2.amazonaws.com/123456789012/MyQueueName"), + } + _, err := sqsClient.SendMessage(context.Background(), sendMessageInput) + require.NoError(t, err) + + spans := mt.FinishedSpans() + require.Len(t, spans, 1) + + s := spans[0] + assert.Equal(t, "SQS.request", s.OperationName()) + assert.Equal(t, "SendMessage", s.Tag("aws.operation")) + assert.Equal(t, "SQS", s.Tag("aws.service")) + assert.Equal(t, "MyQueueName", s.Tag("queuename")) + assert.Equal(t, "SQS.SendMessage", s.Tag(ext.ResourceName)) + assert.Equal(t, "aws.SQS", s.Tag(ext.ServiceName)) + + // Check for trace context injection + assert.NotNil(t, sendMessageInput.MessageAttributes) + assert.Contains(t, sendMessageInput.MessageAttributes, "_datadog") + ddAttr := sendMessageInput.MessageAttributes["_datadog"] + assert.Equal(t, "String", *ddAttr.DataType) + assert.NotEmpty(t, *ddAttr.StringValue) + + // Decode and verify the injected trace context + var traceContext map[string]string + err = json.Unmarshal([]byte(*ddAttr.StringValue), &traceContext) + assert.NoError(t, err) + assert.Contains(t, traceContext, "x-datadog-trace-id") + assert.Contains(t, traceContext, "x-datadog-parent-id") + assert.NotEmpty(t, traceContext["x-datadog-trace-id"]) + assert.NotEmpty(t, traceContext["x-datadog-parent-id"]) +} + func TestAppendMiddlewareS3ListObjects(t *testing.T) { tests := []struct { name string @@ -441,6 +503,22 @@ func TestAppendMiddlewareSnsPublish(t *testing.T) { assert.Equal(t, server.URL+"/", s.Tag(ext.HTTPURL)) assert.Equal(t, "aws/aws-sdk-go-v2/aws", s.Tag(ext.Component)) assert.Equal(t, ext.SpanKindClient, s.Tag(ext.SpanKind)) + + // Check for trace context injection + assert.NotNil(t, tt.publishInput.MessageAttributes) + assert.Contains(t, tt.publishInput.MessageAttributes, "_datadog") + ddAttr := tt.publishInput.MessageAttributes["_datadog"] + assert.Equal(t, "String", *ddAttr.DataType) + assert.NotEmpty(t, *ddAttr.StringValue) + + // Decode and verify the injected trace context + var traceContext map[string]string + err := json.Unmarshal([]byte(*ddAttr.StringValue), &traceContext) + assert.NoError(t, err) + assert.Contains(t, traceContext, "x-datadog-trace-id") + assert.Contains(t, traceContext, "x-datadog-parent-id") + assert.NotEmpty(t, traceContext["x-datadog-trace-id"]) + assert.NotEmpty(t, traceContext["x-datadog-parent-id"]) }) } } @@ -657,6 +735,62 @@ func TestAppendMiddlewareEventBridgePutRule(t *testing.T) { } } +func TestAppendMiddlewareEventBridgePutEvents(t *testing.T) { + mt := mocktracer.Start() + defer mt.Stop() + + expectedStatusCode := 200 + server := mockAWS(expectedStatusCode) + defer server.Close() + + resolver := aws.EndpointResolverFunc(func(service, region string) (aws.Endpoint, error) { + return aws.Endpoint{ + PartitionID: "aws", + URL: server.URL, + SigningRegion: "eu-west-1", + }, nil + }) + + awsCfg := aws.Config{ + Region: "eu-west-1", + Credentials: aws.AnonymousCredentials{}, + EndpointResolver: resolver, + } + + AppendMiddleware(&awsCfg) + + eventbridgeClient := eventbridge.NewFromConfig(awsCfg) + putEventsInput := &eventbridge.PutEventsInput{ + Entries: []eventBridgeTypes.PutEventsRequestEntry{ + { + EventBusName: aws.String("my-event-bus"), + Detail: aws.String(`{"key": "value"}`), + }, + }, + } + eventbridgeClient.PutEvents(context.Background(), putEventsInput) + + spans := mt.FinishedSpans() + require.Len(t, spans, 1) + + s := spans[0] + assert.Equal(t, "PutEvents", s.Tag("aws.operation")) + assert.Equal(t, "EventBridge.PutEvents", s.Tag(ext.ResourceName)) + + // Check for trace context injection + assert.Len(t, putEventsInput.Entries, 1) + entry := putEventsInput.Entries[0] + var detail map[string]interface{} + err := json.Unmarshal([]byte(*entry.Detail), &detail) + assert.NoError(t, err) + assert.Contains(t, detail, "_datadog") + ddData, ok := detail["_datadog"].(map[string]interface{}) + assert.True(t, ok) + assert.Contains(t, ddData, "x-datadog-start-time") + assert.Contains(t, ddData, "x-datadog-resource-name") + assert.Equal(t, "my-event-bus", ddData["x-datadog-resource-name"]) +} + func TestAppendMiddlewareSfnDescribeStateMachine(t *testing.T) { tests := []struct { name string @@ -971,8 +1105,8 @@ func TestMessagingNamingSchema(t *testing.T) { _, err = sqsClient.SendMessage(ctx, msg) require.NoError(t, err) - entry := types.SendMessageBatchRequestEntry{Id: aws.String("1"), MessageBody: aws.String("body")} - batchMsg := &sqs.SendMessageBatchInput{QueueUrl: sqsResp.QueueUrl, Entries: []types.SendMessageBatchRequestEntry{entry}} + entry := sqsTypes.SendMessageBatchRequestEntry{Id: aws.String("1"), MessageBody: aws.String("body")} + batchMsg := &sqs.SendMessageBatchInput{QueueUrl: sqsResp.QueueUrl, Entries: []sqsTypes.SendMessageBatchRequestEntry{entry}} _, err = sqsClient.SendMessageBatch(ctx, batchMsg) require.NoError(t, err) diff --git a/contrib/aws/aws-sdk-go-v2/aws/eventbridge/eventbridge_test.go b/contrib/aws/aws-sdk-go-v2/aws/eventbridge/eventbridge_test.go new file mode 100644 index 0000000000..5414799476 --- /dev/null +++ b/contrib/aws/aws-sdk-go-v2/aws/eventbridge/eventbridge_test.go @@ -0,0 +1,156 @@ +package eventbridge + +import ( + "context" + "encoding/json" + "strconv" + "testing" + + "github.com/aws/aws-sdk-go-v2/aws" + "github.com/aws/aws-sdk-go-v2/service/eventbridge" + "github.com/aws/aws-sdk-go-v2/service/eventbridge/types" + "github.com/aws/smithy-go/middleware" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "gopkg.in/DataDog/dd-trace-go.v1/ddtrace/mocktracer" + "gopkg.in/DataDog/dd-trace-go.v1/ddtrace/tracer" +) + +type testCarrier struct { + m map[string]string +} + +func (c *testCarrier) Set(key, val string) { + c.m[key] = val +} + +func (c *testCarrier) ForeachKey(handler func(key, val string) error) error { + for k, v := range c.m { + if err := handler(k, v); err != nil { + return err + } + } + return nil +} + +func TestEnrichOperation(t *testing.T) { + mt := mocktracer.Start() + defer mt.Stop() + + ctx := context.Background() + _, ctx = tracer.StartSpanFromContext(ctx, "test-span") + + input := middleware.InitializeInput{ + Parameters: &eventbridge.PutEventsInput{ + Entries: []types.PutEventsRequestEntry{ + { + Detail: aws.String(`{"key": "value"}`), + EventBusName: aws.String("test-bus"), + }, + { + Detail: aws.String(`{"another": "data"}`), + EventBusName: aws.String("test-bus-2"), + }, + }, + }, + } + + EnrichOperation(ctx, input, "PutEvents") + + params, ok := input.Parameters.(*eventbridge.PutEventsInput) + require.True(t, ok) + require.Len(t, params.Entries, 2) + + for _, entry := range params.Entries { + var detail map[string]interface{} + err := json.Unmarshal([]byte(*entry.Detail), &detail) + require.NoError(t, err) + + assert.Contains(t, detail, datadogKey) + ddData, ok := detail[datadogKey].(map[string]interface{}) + require.True(t, ok) + + assert.Contains(t, ddData, startTimeKey) + assert.Contains(t, ddData, resourceNameKey) + assert.Equal(t, *entry.EventBusName, ddData[resourceNameKey]) + } +} + +func TestInjectTraceContext(t *testing.T) { + mt := mocktracer.Start() + defer mt.Stop() + + ctx := context.Background() + span, ctx := tracer.StartSpanFromContext(ctx, "test-span") + + tests := []struct { + name string + entry types.PutEventsRequestEntry + expected func(*testing.T, *types.PutEventsRequestEntry) + }{ + { + name: "Inject into empty detail", + entry: types.PutEventsRequestEntry{ + EventBusName: aws.String("test-bus"), + }, + expected: func(t *testing.T, entry *types.PutEventsRequestEntry) { + assert.NotNil(t, entry.Detail) + var detail map[string]interface{} + err := json.Unmarshal([]byte(*entry.Detail), &detail) + require.NoError(t, err) + assert.Contains(t, detail, datadogKey) + }, + }, + { + name: "Inject into existing detail", + entry: types.PutEventsRequestEntry{ + Detail: aws.String(`{"existing": "data"}`), + EventBusName: aws.String("test-bus"), + }, + expected: func(t *testing.T, entry *types.PutEventsRequestEntry) { + var detail map[string]interface{} + err := json.Unmarshal([]byte(*entry.Detail), &detail) + require.NoError(t, err) + assert.Contains(t, detail, "existing") + assert.Equal(t, "data", detail["existing"]) + assert.Contains(t, detail, datadogKey) + }, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + injectTraceContext(ctx, &tt.entry) + tt.expected(t, &tt.entry) + + var detail map[string]interface{} + err := json.Unmarshal([]byte(*tt.entry.Detail), &detail) + require.NoError(t, err) + + ddData := detail[datadogKey].(map[string]interface{}) + assert.Contains(t, ddData, startTimeKey) + assert.Contains(t, ddData, resourceNameKey) + assert.Equal(t, *tt.entry.EventBusName, ddData[resourceNameKey]) + + // Check that start time exists and is not empty + startTimeStr, ok := ddData[startTimeKey].(string) + assert.True(t, ok) + startTime, err := strconv.ParseInt(startTimeStr, 10, 64) + assert.NoError(t, err) + assert.Greater(t, startTime, int64(0)) + + var carrier testCarrier + carrier.m = make(map[string]string) + for k, v := range ddData { + if s, ok := v.(string); ok { + carrier.m[k] = s + } + } + + extractedSpanContext, err := tracer.Extract(&carrier) + assert.NoError(t, err) + assert.Equal(t, span.Context().TraceID(), extractedSpanContext.TraceID()) + assert.Equal(t, span.Context().SpanID(), extractedSpanContext.SpanID()) + }) + } +} diff --git a/contrib/aws/aws-sdk-go-v2/aws/sns/sns_test.go b/contrib/aws/aws-sdk-go-v2/aws/sns/sns_test.go new file mode 100644 index 0000000000..9c04097d0b --- /dev/null +++ b/contrib/aws/aws-sdk-go-v2/aws/sns/sns_test.go @@ -0,0 +1,191 @@ +package sns + +import ( + "context" + "encoding/json" + "fmt" + "testing" + + "github.com/aws/aws-sdk-go-v2/aws" + "github.com/aws/aws-sdk-go-v2/service/sns" + "github.com/aws/aws-sdk-go-v2/service/sns/types" + "github.com/aws/smithy-go/middleware" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "gopkg.in/DataDog/dd-trace-go.v1/ddtrace/mocktracer" + "gopkg.in/DataDog/dd-trace-go.v1/ddtrace/tracer" +) + +type testCarrier struct { + m map[string]string +} + +func (c *testCarrier) Set(key, val string) { + c.m[key] = val +} + +func (c *testCarrier) ForeachKey(handler func(key, val string) error) error { + for k, v := range c.m { + if err := handler(k, v); err != nil { + return err + } + } + return nil +} + +func TestEnrichOperation(t *testing.T) { + tests := []struct { + name string + operation string + input middleware.InitializeInput + setup func(context.Context) context.Context + check func(*testing.T, middleware.InitializeInput) + }{ + { + name: "Publish", + operation: "Publish", + input: middleware.InitializeInput{ + Parameters: &sns.PublishInput{ + Message: aws.String("test message"), + TopicArn: aws.String("arn:aws:sns:us-east-1:123456789012:test-topic"), + }, + }, + setup: func(ctx context.Context) context.Context { + _, ctx = tracer.StartSpanFromContext(ctx, "test-span") + return ctx + }, + check: func(t *testing.T, in middleware.InitializeInput) { + params, ok := in.Parameters.(*sns.PublishInput) + require.True(t, ok) + require.NotNil(t, params) + require.NotNil(t, params.MessageAttributes) + assert.Contains(t, params.MessageAttributes, datadogKey) + assert.NotNil(t, params.MessageAttributes[datadogKey].DataType) + assert.Equal(t, "String", *params.MessageAttributes[datadogKey].DataType) + assert.NotNil(t, params.MessageAttributes[datadogKey].StringValue) + assert.NotEmpty(t, *params.MessageAttributes[datadogKey].StringValue) + }, + }, + { + name: "PublishBatch", + operation: "PublishBatch", + input: middleware.InitializeInput{ + Parameters: &sns.PublishBatchInput{ + TopicArn: aws.String("arn:aws:sns:us-east-1:123456789012:test-topic"), + PublishBatchRequestEntries: []types.PublishBatchRequestEntry{ + { + Id: aws.String("1"), + Message: aws.String("test message 1"), + }, + { + Id: aws.String("2"), + Message: aws.String("test message 2"), + }, + }, + }, + }, + setup: func(ctx context.Context) context.Context { + _, ctx = tracer.StartSpanFromContext(ctx, "test-span") + return ctx + }, + check: func(t *testing.T, in middleware.InitializeInput) { + params, ok := in.Parameters.(*sns.PublishBatchInput) + require.True(t, ok) + require.NotNil(t, params) + require.NotNil(t, params.PublishBatchRequestEntries) + require.Len(t, params.PublishBatchRequestEntries, 2) + + for _, entry := range params.PublishBatchRequestEntries { + require.NotNil(t, entry.MessageAttributes) + assert.Contains(t, entry.MessageAttributes, datadogKey) + assert.NotNil(t, entry.MessageAttributes[datadogKey].DataType) + assert.Equal(t, "String", *entry.MessageAttributes[datadogKey].DataType) + assert.NotNil(t, entry.MessageAttributes[datadogKey].StringValue) + assert.NotEmpty(t, *entry.MessageAttributes[datadogKey].StringValue) + } + }, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + mt := mocktracer.Start() + defer mt.Stop() + + ctx := context.Background() + if tt.setup != nil { + ctx = tt.setup(ctx) + } + + EnrichOperation(ctx, tt.input, tt.operation) + + if tt.check != nil { + tt.check(t, tt.input) + } + }) + } +} + +func TestInjectTraceContext(t *testing.T) { + tests := []struct { + name string + existingAttributes int + expectInjection bool + }{ + { + name: "Inject with no existing attributes", + existingAttributes: 0, + expectInjection: true, + }, + { + name: "Inject with some existing attributes", + existingAttributes: 5, + expectInjection: true, + }, + { + name: "No injection when at max attributes", + existingAttributes: maxMessageAttributes, + expectInjection: false, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + mt := mocktracer.Start() + defer mt.Stop() + + ctx := context.Background() + span, ctx := tracer.StartSpanFromContext(ctx, "test-span") + + messageAttributes := make(map[string]types.MessageAttributeValue) + for i := 0; i < tt.existingAttributes; i++ { + messageAttributes[fmt.Sprintf("attr%d", i)] = types.MessageAttributeValue{ + DataType: aws.String("String"), + StringValue: aws.String("value"), + } + } + + injectTraceContext(ctx, messageAttributes) + + if tt.expectInjection { + assert.Contains(t, messageAttributes, datadogKey) + assert.NotNil(t, messageAttributes[datadogKey].DataType) + assert.Equal(t, "String", *messageAttributes[datadogKey].DataType) + assert.NotNil(t, messageAttributes[datadogKey].StringValue) + assert.NotEmpty(t, *messageAttributes[datadogKey].StringValue) + + var carrier testCarrier + carrier.m = make(map[string]string) + err := json.Unmarshal([]byte(*messageAttributes[datadogKey].StringValue), &carrier.m) + assert.NoError(t, err) + + extractedSpanContext, err := tracer.Extract(&carrier) + assert.NoError(t, err) + assert.Equal(t, span.Context().TraceID(), extractedSpanContext.TraceID()) + assert.Equal(t, span.Context().SpanID(), extractedSpanContext.SpanID()) + } else { + assert.NotContains(t, messageAttributes, datadogKey) + } + }) + } +} diff --git a/contrib/aws/aws-sdk-go-v2/aws/sqs/sqs_test.go b/contrib/aws/aws-sdk-go-v2/aws/sqs/sqs_test.go new file mode 100644 index 0000000000..559a904f21 --- /dev/null +++ b/contrib/aws/aws-sdk-go-v2/aws/sqs/sqs_test.go @@ -0,0 +1,194 @@ +package sqs + +import ( + "context" + "encoding/json" + "fmt" + "github.com/aws/aws-sdk-go-v2/service/sqs/types" + "testing" + + "github.com/aws/aws-sdk-go-v2/aws" + "github.com/aws/aws-sdk-go-v2/service/sqs" + "github.com/aws/smithy-go/middleware" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "gopkg.in/DataDog/dd-trace-go.v1/ddtrace/mocktracer" + "gopkg.in/DataDog/dd-trace-go.v1/ddtrace/tracer" +) + +type testCarrier struct { + m map[string]string +} + +func (c *testCarrier) Set(key, val string) { + c.m[key] = val +} + +func (c *testCarrier) ForeachKey(handler func(key, val string) error) error { + for k, v := range c.m { + if err := handler(k, v); err != nil { + return err + } + } + return nil +} + +func TestEnrichOperation(t *testing.T) { + tests := []struct { + name string + operation string + input middleware.InitializeInput + setup func(context.Context) context.Context + check func(*testing.T, middleware.InitializeInput) + }{ + { + name: "SendMessage", + operation: "SendMessage", + input: middleware.InitializeInput{ + Parameters: &sqs.SendMessageInput{ + MessageBody: aws.String("test message"), + QueueUrl: aws.String("https://sqs.us-east-1.amazonaws.com/1234567890/test-queue"), + }, + }, + setup: func(ctx context.Context) context.Context { + _, ctx = tracer.StartSpanFromContext(ctx, "test-span") + return ctx + }, + check: func(t *testing.T, in middleware.InitializeInput) { + params, ok := in.Parameters.(*sqs.SendMessageInput) + require.True(t, ok) + require.NotNil(t, params) + require.NotNil(t, params.MessageAttributes) + assert.Contains(t, params.MessageAttributes, datadogKey) + assert.NotNil(t, params.MessageAttributes[datadogKey].DataType) + assert.Equal(t, "String", *params.MessageAttributes[datadogKey].DataType) + assert.NotNil(t, params.MessageAttributes[datadogKey].StringValue) + assert.NotEmpty(t, *params.MessageAttributes[datadogKey].StringValue) + }, + }, + { + name: "SendMessageBatch", + operation: "SendMessageBatch", + input: middleware.InitializeInput{ + Parameters: &sqs.SendMessageBatchInput{ + QueueUrl: aws.String("https://sqs.us-east-1.amazonaws.com/1234567890/test-queue"), + Entries: []types.SendMessageBatchRequestEntry{ + { + Id: aws.String("1"), + MessageBody: aws.String("test message 1"), + }, + { + Id: aws.String("2"), + MessageBody: aws.String("test message 2"), + }, + { + Id: aws.String("3"), + MessageBody: aws.String("test message 3"), + }, + }, + }, + }, + setup: func(ctx context.Context) context.Context { + _, ctx = tracer.StartSpanFromContext(ctx, "test-span") + return ctx + }, + check: func(t *testing.T, in middleware.InitializeInput) { + params, ok := in.Parameters.(*sqs.SendMessageBatchInput) + require.True(t, ok) + require.NotNil(t, params) + require.NotNil(t, params.Entries) + require.Len(t, params.Entries, 3) + + for _, entry := range params.Entries { + require.NotNil(t, entry.MessageAttributes) + assert.Contains(t, entry.MessageAttributes, datadogKey) + assert.NotNil(t, entry.MessageAttributes[datadogKey].DataType) + assert.Equal(t, "String", *entry.MessageAttributes[datadogKey].DataType) + assert.NotNil(t, entry.MessageAttributes[datadogKey].StringValue) + assert.NotEmpty(t, *entry.MessageAttributes[datadogKey].StringValue) + } + }, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + mt := mocktracer.Start() + defer mt.Stop() + + ctx := context.Background() + if tt.setup != nil { + ctx = tt.setup(ctx) + } + + EnrichOperation(ctx, tt.input, tt.operation) + + if tt.check != nil { + tt.check(t, tt.input) + } + }) + } +} + +func TestInjectTraceContext(t *testing.T) { + tests := []struct { + name string + existingAttributes int + expectInjection bool + }{ + { + name: "Inject with no existing attributes", + existingAttributes: 0, + expectInjection: true, + }, + { + name: "Inject with some existing attributes", + existingAttributes: 5, + expectInjection: true, + }, + { + name: "No injection when at max attributes", + existingAttributes: maxMessageAttributes, + expectInjection: false, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + mt := mocktracer.Start() + defer mt.Stop() + + ctx := context.Background() + span, ctx := tracer.StartSpanFromContext(ctx, "test-span") + + messageAttributes := make(map[string]types.MessageAttributeValue) + for i := 0; i < tt.existingAttributes; i++ { + messageAttributes[fmt.Sprintf("attr%d", i)] = types.MessageAttributeValue{ + DataType: aws.String("String"), + StringValue: aws.String("value"), + } + } + + injectTraceContext(ctx, messageAttributes) + + if tt.expectInjection { + assert.Contains(t, messageAttributes, datadogKey) + assert.NotNil(t, messageAttributes[datadogKey].DataType) + assert.Equal(t, "String", *messageAttributes[datadogKey].DataType) + assert.NotNil(t, messageAttributes[datadogKey].StringValue) + assert.NotEmpty(t, *messageAttributes[datadogKey].StringValue) + + var carrier testCarrier + err := json.Unmarshal([]byte(*messageAttributes[datadogKey].StringValue), &carrier.m) + assert.NoError(t, err) + + extractedSpanContext, err := tracer.Extract(&carrier) + assert.NoError(t, err) + assert.Equal(t, span.Context().TraceID(), extractedSpanContext.TraceID()) + assert.Equal(t, span.Context().SpanID(), extractedSpanContext.SpanID()) + } else { + assert.NotContains(t, messageAttributes, datadogKey) + } + }) + } +}