Move trace details propagated when message attributes are readonly to SnsInterceptorTest

Author: ygree

dd-java-agent/instrumentation/aws-java/aws-java-sns-2.0/src/test/groovy/SnsClientTest.groovy

@@ -13,7 +13,6 @@ import software.amazon.awssdk.auth.credentials.AwsBasicCredentials
 import software.amazon.awssdk.auth.credentials.StaticCredentialsProvider
 import software.amazon.awssdk.regions.Region
 import software.amazon.awssdk.services.sns.SnsClient
-import software.amazon.awssdk.services.sns.model.MessageAttributeValue
 import software.amazon.awssdk.services.sns.model.PublishResponse
 import software.amazon.awssdk.services.sqs.SqsClient
 import software.amazon.awssdk.services.sqs.model.QueueAttributeName
@@ -41,15 +40,15 @@ abstract class SnsClientTest extends VersionedNamingTestBase {
     LOCALSTACK.start()
     def endPoint = "http://" + LOCALSTACK.getHost() + ":" + LOCALSTACK.getMappedPort(4566)
     snsClient = SnsClient.builder()
-    .endpointOverride(URI.create(endPoint))
-    .region(Region.of("us-east-1"))
-    .credentialsProvider(StaticCredentialsProvider.create(AwsBasicCredentials.create("test", "test")))
-    .build()
+      .endpointOverride(URI.create(endPoint))
+      .region(Region.of("us-east-1"))
+      .credentialsProvider(StaticCredentialsProvider.create(AwsBasicCredentials.create("test", "test")))
+      .build()
     sqsClient = SqsClient.builder()
-    .endpointOverride(URI.create(endPoint))
-    .region(Region.of("us-east-1"))
-    .credentialsProvider(StaticCredentialsProvider.create(AwsBasicCredentials.create("test", "test")))
-    .build()
+      .endpointOverride(URI.create(endPoint))
+      .region(Region.of("us-east-1"))
+      .credentialsProvider(StaticCredentialsProvider.create(AwsBasicCredentials.create("test", "test")))
+      .build()
     testQueueURL = sqsClient.createQueue {  it.queueName("testqueue") }.queueUrl()
     testQueueARN = sqsClient.getQueueAttributes {it.queueUrl(testQueueURL).attributeNames(QueueAttributeName.QUEUE_ARN)}.attributes().get(QueueAttributeName.QUEUE_ARN)
     testTopicARN = snsClient.createTopic { it.name("testtopic") }.topicArn()
@@ -82,30 +81,6 @@ abstract class SnsClientTest extends VersionedNamingTestBase {
   abstract String expectedOperation(String awsService, String awsOperation)
   abstract String expectedService(String awsService, String awsOperation)
 
-  def "trace details propagated when message attributes are readonly"() {
-    when:
-    TEST_WRITER.clear()
-
-    def headers = new HashMap<String, MessageAttributeValue>()
-    headers.put("mykey", MessageAttributeValue.builder().stringValue("myvalue").dataType("String").build())
-    def readonlyHeaders = Collections.unmodifiableMap(headers)
-    snsClient.publish(b -> b.message("sometext").topicArn(testTopicARN).messageAttributes(readonlyHeaders))
-
-    def message = sqsClient.receiveMessage {
-      it.queueUrl(testQueueURL).waitTimeSeconds(3)
-    }.messages().get(0)
-
-    def messageBody = new JsonSlurper().parseText(message.body())
-
-    then:
-    // injected value is here
-    String injectedValue = messageBody["MessageAttributes"]["_datadog"]["Value"]
-    injectedValue.length() > 0
-
-    // original header value is still present
-    messageBody["MessageAttributes"]["mykey"] != null
-  }
-
   def "trace details propagated via SNS system message attributes"() {
     when:
     TEST_WRITER.clear()
@@ -214,7 +189,7 @@ abstract class SnsClientTest extends VersionedNamingTestBase {
     TEST_WRITER.clear()
     snsClient.publish { req ->
       req.message("test message")
-      .topicArn(testTopicARN)
+        .topicArn(testTopicARN)
     }
 
     def message = sqsClient.receiveMessage { it.queueUrl(testQueueURL).waitTimeSeconds(3) }.messages().get(0)
@@ -339,4 +314,3 @@ class SnsClientV1DataStreamsForkedTest extends SnsClientTest {
     1
   }
 }
-

dd-java-agent/instrumentation/aws-java/aws-java-sns-2.0/src/test/java/datadog/trace/instrumentation/aws/v2/sns/SnsInterceptorTest.java

@@ -2,10 +2,13 @@
 
 import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.junit.jupiter.api.Assertions.assertFalse;
+import static org.junit.jupiter.api.Assertions.assertNotSame;
 import static org.junit.jupiter.api.Assertions.assertTrue;
 
 import datadog.context.Context;
 import java.util.Arrays;
+import java.util.Collections;
+import java.util.HashMap;
 import java.util.LinkedHashMap;
 import java.util.Map;
 import java.util.stream.Collectors;
@@ -14,6 +17,7 @@
 import software.amazon.awssdk.services.sns.model.MessageAttributeValue;
 import software.amazon.awssdk.services.sns.model.PublishBatchRequest;
 import software.amazon.awssdk.services.sns.model.PublishBatchRequestEntry;
+import software.amazon.awssdk.services.sns.model.PublishRequest;
 
 public class SnsInterceptorTest {
 
@@ -65,4 +69,27 @@ private static Map<String, MessageAttributeValue> stringAttributes(int count) {
     }
     return attributes;
   }
+
+  @Test
+  void publishPreservesReadonlyAttributesWhileAddingDatadogContext() {
+    Map<String, MessageAttributeValue> headers = new HashMap<>();
+    headers.put(
+        "mykey", MessageAttributeValue.builder().dataType("String").stringValue("myvalue").build());
+    Map<String, MessageAttributeValue> readonlyHeaders = Collections.unmodifiableMap(headers);
+
+    PublishRequest request =
+        PublishRequest.builder()
+            .topicArn("arn:aws:sns:us-east-1:123456789012:test-topic")
+            .message("sometext")
+            .messageAttributes(readonlyHeaders)
+            .build();
+
+    PublishRequest modified =
+        (PublishRequest) new SnsInterceptor().modifyRequest(() -> request, executionAttributes());
+
+    assertNotSame(readonlyHeaders, modified.messageAttributes());
+    assertEquals("myvalue", modified.messageAttributes().get("mykey").stringValue());
+    assertTrue(modified.messageAttributes().containsKey("_datadog"));
+    assertFalse(readonlyHeaders.containsKey("_datadog"));
+  }
 }