From 80088f926286bdb6b39eff641ef4e978cf23541f Mon Sep 17 00:00:00 2001
From: Oren Ben-Meir <orenbm11@gmail.com>
Date: Thu, 16 Jan 2025 12:52:44 -0500
Subject: [PATCH] Test DT headers in V1 SDK

---
 .../amazonaws/services/sqs/SqsClientTest.java | 43 ++++++++++++++++---
 1 file changed, 37 insertions(+), 6 deletions(-)

diff --git a/instrumentation/aws-java-sdk-sqs-1.10.44/src/test/java/com/amazonaws/services/sqs/SqsClientTest.java b/instrumentation/aws-java-sdk-sqs-1.10.44/src/test/java/com/amazonaws/services/sqs/SqsClientTest.java
index 0c037b731c..c5910e5ba6 100644
--- a/instrumentation/aws-java-sdk-sqs-1.10.44/src/test/java/com/amazonaws/services/sqs/SqsClientTest.java
+++ b/instrumentation/aws-java-sdk-sqs-1.10.44/src/test/java/com/amazonaws/services/sqs/SqsClientTest.java
@@ -13,12 +13,14 @@
 import com.amazonaws.services.sqs.model.CreateQueueResult;
 import com.amazonaws.services.sqs.model.ReceiveMessageRequest;
 import com.amazonaws.services.sqs.model.SendMessageBatchRequest;
+import com.amazonaws.services.sqs.model.SendMessageBatchRequestEntry;
 import com.amazonaws.services.sqs.model.SendMessageRequest;
 import com.newrelic.agent.introspec.InstrumentationTestConfig;
 import com.newrelic.agent.introspec.InstrumentationTestRunner;
 import com.newrelic.agent.introspec.Introspector;
 import com.newrelic.agent.introspec.TracedMetricData;
 import com.newrelic.api.agent.Trace;
+import com.newrelic.utils.SqsV1Util;
 import org.elasticmq.NodeAddress;
 import org.elasticmq.rest.sqs.SQSRestServer;
 import org.elasticmq.rest.sqs.SQSRestServerBuilder;
@@ -27,7 +29,13 @@
 import org.junit.Test;
 import org.junit.runner.RunWith;
 
+import java.util.Arrays;
+import java.util.HashSet;
+import java.util.Set;
+
 import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertTrue;
 
 @RunWith(InstrumentationTestRunner.class)
 @InstrumentationTestConfig(includePrefixes = { "com.amazonaws.services.sqs" }, configName = "dt_enabled.yml")
@@ -67,7 +75,13 @@ public static void afterClass() {
     @Test
     public void testSendMessage() {
         Introspector introspector = InstrumentationTestRunner.getIntrospector();
-        sendMessageRequest();
+        SendMessageRequest request = sendMessageRequest();
+
+        Set<String> dtHeaders = new HashSet<>(Arrays.asList(SqsV1Util.DT_HEADERS));
+        boolean containsDtHeaders = request.getMessageAttributes().entrySet().stream().anyMatch(e ->
+                dtHeaders.contains(e.getKey()) && e.getValue() != null && !e.getValue().getStringValue().isEmpty());
+        assertTrue("Message request must contain headers", containsDtHeaders);
+
         assertEquals(1, introspector.getFinishedTransactionCount(10000));
 
         String txName = introspector.getTransactionNames().iterator().next();
@@ -77,9 +91,17 @@ public void testSendMessage() {
     @Test
     public void testSendMessageBatch() {
         Introspector introspector = InstrumentationTestRunner.getIntrospector();
-        sendMessageBatch();
+        SendMessageBatchRequest request = sendMessageBatch();
         assertEquals(1, introspector.getFinishedTransactionCount(10000));
 
+        Set<String> dtHeaders = new HashSet<>(Arrays.asList(SqsV1Util.DT_HEADERS));
+        assertFalse("Batch request must contain at least one entry", request.getEntries().isEmpty());
+        for (SendMessageBatchRequestEntry entry: request.getEntries()) {
+            boolean containsDtHeaders = entry.getMessageAttributes().entrySet().stream().anyMatch(e ->
+                    dtHeaders.contains(e.getKey()) && e.getValue() != null && !e.getValue().getStringValue().isEmpty());
+            assertTrue("Message entry must contain headers", containsDtHeaders);
+        }
+
         String txName = introspector.getTransactionNames().iterator().next();
         checkScopedMetricCount(txName, "MessageBroker/SQS/Queue/Produce/Named/" + QUEUE_NAME, 1);
     }
@@ -95,15 +117,24 @@ public void testReceiveMessage() {
     }
 
     @Trace(dispatcher = true)
-    private void sendMessageRequest() {
+    private SendMessageRequest sendMessageRequest() {
         SendMessageRequest request = (new SendMessageRequest()).withQueueUrl(queueUrl).withMessageBody("body");
         sqsClient.sendMessage(request);
+        return request;
     }
 
     @Trace(dispatcher = true)
-    private void sendMessageBatch() {
-        SendMessageBatchRequest request = (new SendMessageBatchRequest()).withQueueUrl(queueUrl);
-        sqsClient.sendMessageBatch(request);
+    private SendMessageBatchRequest sendMessageBatch() {
+        SendMessageBatchRequestEntry entry = new SendMessageBatchRequestEntry();
+        SendMessageBatchRequest request = (new SendMessageBatchRequest()).withQueueUrl(queueUrl)
+                .withEntries(entry);
+        try {
+            sqsClient.sendMessageBatch(request);
+        } catch (Exception e) {
+            // Do nothing
+        }
+
+        return request;
     }
 
     @Trace(dispatcher = true)