Fix KafkaStreamsTest

Author: smola

dd-java-agent/instrumentation/kafka-streams-0.11/src/latestDepTest/groovy/KafkaStreamsTest.groovy

@@ -27,6 +27,7 @@ import spock.lang.Shared
 
 import java.util.concurrent.LinkedBlockingQueue
 import java.util.concurrent.TimeUnit
+import java.util.regex.Pattern
 
 @Flaky("https://github.com/DataDog/dd-trace-java/issues/3865")
 class KafkaStreamsTest extends AgentTestRunner {
@@ -39,6 +40,11 @@ class KafkaStreamsTest extends AgentTestRunner {
   @Shared
   EmbeddedKafkaBroker embeddedKafka = kafkaRule.embeddedKafka
 
+  def setup() {
+    // Filter out additional traces for kafka.poll operation, otherwise, there will be more traces than expected.
+    TEST_WRITER.setFilter { trace -> trace[0].operationName.toString() != 'kafka.poll' }
+  }
+
   @Override
   protected boolean isDataStreamsEnabled() {
     return true
@@ -135,6 +141,7 @@ class KafkaStreamsTest extends AgentTestRunner {
             if ({ isDataStreamsEnabled()}) {
               "$DDTags.PATHWAY_HASH" { String }
             }
+            "$InstrumentationTags.KAFKA_BOOTSTRAP_SERVERS" Pattern.compile("127.0.0.1:[0-9]+")
             defaultTagsNoPeerService()
           }
         }
@@ -159,7 +166,6 @@ class KafkaStreamsTest extends AgentTestRunner {
             "$InstrumentationTags.PARTITION" { it >= 0 }
             "$InstrumentationTags.OFFSET" 0
             "$InstrumentationTags.PROCESSOR_NAME" "KSTREAM-SOURCE-0000000000"
-            "$InstrumentationTags.MESSAGING_DESTINATION_NAME" "$STREAM_PENDING"
             "asdf" "testing"
             if ({ isDataStreamsEnabled()}) {
               "$DDTags.PATHWAY_HASH" { String }
@@ -186,6 +192,7 @@ class KafkaStreamsTest extends AgentTestRunner {
             if ({ isDataStreamsEnabled()}) {
               "$DDTags.PATHWAY_HASH" { String }
             }
+            "$InstrumentationTags.KAFKA_BOOTSTRAP_SERVERS" Pattern.compile("127.0.0.1:[0-9]+")
             defaultTagsNoPeerService()
           }
         }
@@ -212,6 +219,7 @@ class KafkaStreamsTest extends AgentTestRunner {
             if ({ isDataStreamsEnabled()}) {
               "$DDTags.PATHWAY_HASH" { String }
             }
+            "$InstrumentationTags.KAFKA_BOOTSTRAP_SERVERS" Pattern.compile("127.0.0.1:[0-9]+")
             defaultTags(true)
           }
         }
@@ -226,8 +234,11 @@ class KafkaStreamsTest extends AgentTestRunner {
     if (isDataStreamsEnabled()) {
       StatsGroup originProducerPoint = TEST_DATA_STREAMS_WRITER.groups.find { it.parentHash == 0 }
       verifyAll(originProducerPoint) {
-        edgeTags == ["direction:out", "topic:$STREAM_PENDING", "type:kafka"]
-        edgeTags.size() == 3
+        edgeTags.any { it.startsWith("kafka_cluster_id:") }
+        for (String tag : ["direction:out", "topic:$STREAM_PENDING", "type:kafka"]) {
+          assert edgeTags.contains(tag)
+        }
+        edgeTags.size() == 4
       }
 
       StatsGroup kafkaStreamsConsumerPoint = TEST_DATA_STREAMS_WRITER.groups.find { it.parentHash == originProducerPoint.hash }
@@ -243,14 +254,20 @@ class KafkaStreamsTest extends AgentTestRunner {
 
       StatsGroup kafkaStreamsProducerPoint = TEST_DATA_STREAMS_WRITER.groups.find { it.parentHash == kafkaStreamsConsumerPoint.hash }
       verifyAll(kafkaStreamsProducerPoint) {
-        edgeTags == ["direction:out", "topic:$STREAM_PROCESSED", "type:kafka"]
-        edgeTags.size() == 3
+        edgeTags.any { it.startsWith("kafka_cluster_id:") }
+        for (String tag : ["direction:out", "topic:$STREAM_PROCESSED", "type:kafka"]) {
+          assert edgeTags.contains(tag)
+        }
+        edgeTags.size() == 4
       }
 
       StatsGroup finalConsumerPoint = TEST_DATA_STREAMS_WRITER.groups.find { it.parentHash == kafkaStreamsProducerPoint.hash }
       verifyAll(finalConsumerPoint) {
-        edgeTags == ["direction:in", "group:sender", "topic:$STREAM_PROCESSED".toString(), "type:kafka"]
-        edgeTags.size() == 4
+        edgeTags.any { it.startsWith("kafka_cluster_id:") }
+        for (String tag : ["direction:in", "group:sender", "topic:$STREAM_PROCESSED".toString(), "type:kafka"]) {
+          assert edgeTags.contains(tag)
+        }
+        edgeTags.size() == 5
       }
     }