Test for different headers when multiple propagation styles are set

Author: nhulston

dd-java-agent/instrumentation/aws-java-eventbridge-2.0/src/test/groovy/EventBridgeClientTest.groovy

@@ -27,19 +27,30 @@ class EventBridgeClientTest extends AgentTestRunner {
   .withReuse(true)
   .withStartupTimeout(Duration.ofSeconds(120))
 
-  @Shared SnsClient snsClient
-  @Shared String testTopicARN
-  @Shared String testTopicName
-
-  @Shared EventBridgeClient eventBridgeClient
-  @Shared EventBridgeAsyncClient eventBridgeAsyncClient
-  @Shared String testBusARN
-  @Shared String testBusName
-  @Shared String testRuleName
-
-  @Shared SqsClient sqsClient
-  @Shared String testQueueURL
-  @Shared String testQueueARN
+  @Shared
+  SnsClient snsClient
+  @Shared
+  String testTopicARN
+  @Shared
+  String testTopicName
+
+  @Shared
+  EventBridgeClient eventBridgeClient
+  @Shared
+  EventBridgeAsyncClient eventBridgeAsyncClient
+  @Shared
+  String testBusARN
+  @Shared
+  String testBusName
+  @Shared
+  String testRuleName
+
+  @Shared
+  SqsClient sqsClient
+  @Shared
+  String testQueueURL
+  @Shared
+  String testQueueARN
 
   def setupSpec() {
     LOCALSTACK.start()
@@ -115,6 +126,9 @@ class EventBridgeClientTest extends AgentTestRunner {
     super.configurePreAgent()
     injectSysConfig(GeneralConfig.SERVICE_NAME, "eventbridge")
     injectSysConfig(GeneralConfig.DATA_STREAMS_ENABLED, "true")
+
+    // test propagation styles
+    injectSysConfig('dd.trace.propagation.style', 'datadog,b3single,b3multi,haystack,xray,tracecontext')
   }
 
   def "trace details propagated via EventBridge to SQS (sync)"() {
@@ -443,4 +457,42 @@ class EventBridgeClientTest extends AgentTestRunner {
     assert emptyDetailBody["detail"]["_datadog"]["x-datadog-sent-timestamp"] != null
     assert emptyDetailBody["detail"]["_datadog"]["x-datadog-bus-name"] != null
   }
+
+  def "test propagation styles"() {
+    when:
+    eventBridgeClient.putEvents { req ->
+      req.entries(
+        PutEventsRequestEntry.builder()
+        .source("com.example")
+        .detailType("test")
+        .detail('{"foo":"bar"}')
+        .eventBusName(testBusARN)
+        .build()
+        )
+    }
+
+    def message = sqsClient.receiveMessage { it.queueUrl(testQueueURL).waitTimeSeconds(3) }.messages().get(0)
+    def messageBody = new JsonSlurper().parseText(message.body())
+    def traceContext = messageBody["detail"]["_datadog"]
+
+    then:
+    expectedHeaders.each { header ->
+      assert traceContext[header] != null
+    }
+
+    where:
+    expectedHeaders = [
+      'x-datadog-trace-id',
+      'x-datadog-parent-id',
+      'x-datadog-sampling-priority',
+      'b3',
+      'X-B3-TraceId',
+      'X-B3-SpanId',
+      'Span-ID',
+      'Parent-ID',
+      'X-Amzn-Trace-Id',
+      'traceparent',
+      'tracestate'
+    ]
+  }
 }

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

@@ -65,6 +65,9 @@ abstract class SnsClientTest extends VersionedNamingTestBase {
     // Set a service name that gets sorted early with SORT_BY_NAMES
     injectSysConfig(GeneralConfig.SERVICE_NAME, "A-service")
     injectSysConfig(GeneralConfig.DATA_STREAMS_ENABLED, isDataStreamsEnabled().toString())
+
+    // test propagation styles
+    injectSysConfig('dd.trace.propagation.style', 'datadog,b3single,b3multi,xray,tracecontext')
   }
 
   @Override
@@ -188,6 +191,41 @@ abstract class SnsClientTest extends VersionedNamingTestBase {
     then:
     noExceptionThrown()
   }
+
+  def "test propagation styles"() {
+    when:
+    TEST_WRITER.clear()
+    snsClient.publish { req ->
+      req.message("test message")
+        .topicArn(testTopicARN)
+    }
+
+    def message = sqsClient.receiveMessage { it.queueUrl(testQueueURL).waitTimeSeconds(3) }.messages().get(0)
+    def messageBody = new JsonSlurper().parseText(message.body())
+
+    String base64EncodedString = messageBody["MessageAttributes"]["_datadog"]["Value"]
+    byte[] decodedBytes = base64EncodedString.decodeBase64()
+    String decodedString = new String(decodedBytes, "UTF-8")
+    def traceContext = new JsonSlurper().parseText(decodedString)
+
+    then:
+    expectedHeaders.each { header ->
+      assert traceContext[header] != null, "Header $header is missing"
+    }
+
+    where:
+    expectedHeaders = [
+      'x-datadog-trace-id',
+      'x-datadog-parent-id',
+      'x-datadog-sampling-priority',
+      'b3',
+      'X-B3-TraceId',
+      'X-B3-SpanId',
+      'X-Amzn-Trace-Id',
+      'traceparent',
+      'tracestate'
+    ]
+  }
 }
 
 class SnsClientV0Test extends SnsClientTest {