Using unit testing for trace parent changes

Author: Yoone

dd-java-agent/instrumentation/mongo/common/src/main/java/datadog/trace/instrumentation/mongo/MongoCommentInjector.java

@@ -112,7 +112,7 @@ private static String getHostnameFromEvent(CommandStartedEvent event) {
     return null;
   }
 
-  private static String buildTraceParent(AgentSpan span) {
+  static String buildTraceParent(AgentSpan span) {
     // W3C traceparent format: version-traceId-spanId-flags
     StringBuilder sb = new StringBuilder(2 + 1 + 32 + 1 + 16 + 1 + 2);
     sb.append("00-"); // version

dd-java-agent/instrumentation/mongo/common/src/test/groovy/MongoCommentInjectorTest.groovy

@@ -26,64 +26,36 @@ class MongoCommentInjectorTest extends InstrumentationSpecification {
     comment == null
   }
 
-  def "getComment with full mode builds traceParent in W3C format with sampled flag"() {
+  def "buildTraceParent with sampled flag (SAMPLER_KEEP)"() {
     setup:
-    // Set the propagation mode to full to trigger buildTraceParent call
-    injectSysConfig("dd." + TraceInstrumentationConfig.DB_DBM_PROPAGATION_MODE_MODE, "full")
     def span = TEST_TRACER.buildSpan("test-op").start()
     span.setSamplingPriority(PrioritySampling.SAMPLER_KEEP, 0)
 
-    // Create a mock event
-    def event = Mock(com.mongodb.event.CommandStartedEvent) {
-      getDatabaseName() >> "testdb"
-      getConnectionDescription() >> Mock(com.mongodb.connection.ConnectionDescription) {
-        getServerAddress() >> Mock(com.mongodb.ServerAddress) {
-          getHost() >> "localhost"
-        }
-      }
-    }
-
     when:
-    String comment = MongoCommentInjector.getComment(span, event)
+    String traceParent = MongoCommentInjector.buildTraceParent(span)
 
     then:
-    comment != null
-    comment.contains("dddbs='test-mongo-service'")
-    comment.contains("dde='test'")
-    comment.contains("ddh='localhost'")
-    comment.contains("dddb='testdb'")
-    // W3C traceparent format: 00-{32 hex chars}-{16 hex chars}-01 (sampled)
-    comment ==~ /.*ddtp='00-[0-9a-f]{32}-[0-9a-f]{16}-01'.*/
+    traceParent != null
+    traceParent ==~ /00-[0-9a-f]{32}-[0-9a-f]{16}-01/
+
+    cleanup:
+    span?.finish()
   }
 
-  def "getComment with full mode builds traceParent in W3C format with not sampled flag"() {
+  def "buildTraceParent with not sampled flag (SAMPLER_DROP)"() {
     setup:
-    // Set the propagation mode to full to trigger buildTraceParent call
-    injectSysConfig("dd." + TraceInstrumentationConfig.DB_DBM_PROPAGATION_MODE_MODE, "full")
     def span = TEST_TRACER.buildSpan("test-op").start()
     span.setSamplingPriority(PrioritySampling.SAMPLER_DROP, 0)
 
-    // Create a mock event
-    def event = Mock(com.mongodb.event.CommandStartedEvent) {
-      getDatabaseName() >> "testdb"
-      getConnectionDescription() >> Mock(com.mongodb.connection.ConnectionDescription) {
-        getServerAddress() >> Mock(com.mongodb.ServerAddress) {
-          getHost() >> "localhost"
-        }
-      }
-    }
-
     when:
-    String comment = MongoCommentInjector.getComment(span, event)
+    String traceParent = MongoCommentInjector.buildTraceParent(span)
 
     then:
-    comment != null
-    comment.contains("dddbs='test-mongo-service'")
-    comment.contains("dde='test'")
-    comment.contains("ddh='localhost'")
-    comment.contains("dddb='testdb'")
-    // W3C traceparent format: 00-{32 hex chars}-{16 hex chars}-00 (not sampled)
-    comment ==~ /.*ddtp='00-[0-9a-f]{32}-[0-9a-f]{16}-00'.*/
+    traceParent != null
+    traceParent ==~ /00-[0-9a-f]{32}-[0-9a-f]{16}-00/
+
+    cleanup:
+    span?.finish()
   }
 
   def "injectComment returns null when event is null"() {

dd-java-agent/instrumentation/mongo/driver-3.1/src/test/groovy/MongoJava31ClientTest.groovy

@@ -8,20 +8,15 @@ import com.mongodb.event.CommandFailedEvent
 import com.mongodb.event.CommandListener
 import com.mongodb.event.CommandStartedEvent
 import com.mongodb.event.CommandSucceededEvent
-import datadog.trace.api.sampling.PrioritySampling
-import datadog.trace.bootstrap.instrumentation.api.AgentSpan
 import datadog.trace.core.DDSpan
 import org.bson.BsonDocument
 import org.bson.BsonString
 import org.bson.Document
 import spock.lang.Shared
 
-import java.util.concurrent.ConcurrentLinkedQueue
-
 import static datadog.trace.agent.test.utils.PortUtils.UNUSABLE_PORT
 import static datadog.trace.agent.test.utils.TraceUtils.runUnderTrace
 import static datadog.trace.api.config.TraceInstrumentationConfig.DB_CLIENT_HOST_SPLIT_BY_INSTANCE
-import static datadog.trace.api.config.TraceInstrumentationConfig.DB_DBM_PROPAGATION_MODE_MODE
 import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activeSpan
 
 abstract class MongoJava31ClientTest extends MongoBaseTest {
@@ -236,132 +231,6 @@ abstract class MongoJava31ClientTest extends MongoBaseTest {
     // Unfortunately not caught by our instrumentation.
     assertTraces(0) {}
   }
-
-  def "test DBM propagation with sampled flag"() {
-    setup:
-    injectSysConfig("dd." + DB_DBM_PROPAGATION_MODE_MODE, "full")
-    injectSysConfig("service.name", "test-mongo-service")
-    injectSysConfig("dd.env", "test")
-
-    String collectionName = randomCollectionName()
-    def capturedEvents = new ConcurrentLinkedQueue<CommandStartedEvent>()
-
-    def dbmClient = new MongoClient(new ServerAddress(mongoDbContainer.getHost(), port),
-      MongoClientOptions.builder()
-      .description("dbm-test-client")
-      .addCommandListener(new CommandListener() {
-        @Override
-        void commandStarted(CommandStartedEvent event) {
-          capturedEvents.add(event)
-        }
-
-        @Override
-        void commandSucceeded(CommandSucceededEvent event) {}
-
-        @Override
-        void commandFailed(CommandFailedEvent event) {}
-      })
-      .build())
-
-    DDSpan setupSpan = null
-    MongoCollection<Document> collection = runUnderTrace("setup") {
-      setupSpan = activeSpan() as DDSpan
-      MongoDatabase db = dbmClient.getDatabase(databaseName)
-      db.createCollection(collectionName)
-      return db.getCollection(collectionName)
-    }
-    TEST_WRITER.waitUntilReported(setupSpan)
-    TEST_WRITER.clear()
-    capturedEvents.clear()
-
-    when:
-    runUnderTrace("dbm-test") {
-      AgentSpan span = activeSpan() as AgentSpan
-      span.setSamplingPriority(PrioritySampling.SAMPLER_KEEP, 0)
-      collection.insertOne(new Document("test", "value"))
-    }
-
-    then:
-    // Find the insert command event
-    def insertEvent = capturedEvents.find { it.commandName == "insert" }
-    insertEvent != null
-    def command = insertEvent.command
-    command.containsKey('$comment')
-    def comment = command.getString('$comment').value
-
-    // Verify comment contains DBM tags
-    comment.contains("dddbs='test-mongo-service'")
-    comment.contains("dde='test'")
-
-    // Verify traceparent has sampled flag (01)
-    comment ==~ /.*ddtp='00-[0-9a-f]{32}-[0-9a-f]{16}-01'.*/
-
-    cleanup:
-    dbmClient?.close()
-  }
-
-  def "test DBM propagation with not sampled flag"() {
-    setup:
-    injectSysConfig("dd." + DB_DBM_PROPAGATION_MODE_MODE, "full")
-    injectSysConfig("service.name", "test-mongo-service")
-    injectSysConfig("dd.env", "test")
-
-    String collectionName = randomCollectionName()
-    def capturedEvents = new ConcurrentLinkedQueue<CommandStartedEvent>()
-
-    def dbmClient = new MongoClient(new ServerAddress(mongoDbContainer.getHost(), port),
-      MongoClientOptions.builder()
-      .description("dbm-test-client")
-      .addCommandListener(new CommandListener() {
-        @Override
-        void commandStarted(CommandStartedEvent event) {
-          capturedEvents.add(event)
-        }
-
-        @Override
-        void commandSucceeded(CommandSucceededEvent event) {}
-
-        @Override
-        void commandFailed(CommandFailedEvent event) {}
-      })
-      .build())
-
-    DDSpan setupSpan = null
-    MongoCollection<Document> collection = runUnderTrace("setup") {
-      setupSpan = activeSpan() as DDSpan
-      MongoDatabase db = dbmClient.getDatabase(databaseName)
-      db.createCollection(collectionName)
-      return db.getCollection(collectionName)
-    }
-    TEST_WRITER.waitUntilReported(setupSpan)
-    TEST_WRITER.clear()
-    capturedEvents.clear()
-
-    when:
-    runUnderTrace("dbm-test") {
-      AgentSpan span = activeSpan() as AgentSpan
-      span.setSamplingPriority(PrioritySampling.SAMPLER_DROP, 0)
-      collection.insertOne(new Document("test", "value"))
-    }
-
-    then:
-    // Find the insert command event
-    def insertEvent = capturedEvents.find { it.commandName == "insert" }
-    insertEvent != null
-    def command = insertEvent.command
-    command.containsKey('$comment')
-    def comment = command.getString('$comment').value
-
-    // Verify comment contains DBM tags
-    comment.contains("dddbs='test-mongo-service'")
-    comment.contains("dde='test'")
-
-    // Verify traceparent has not sampled flag (00)
-    comment ==~ /.*ddtp='00-[0-9a-f]{32}-[0-9a-f]{16}-00'.*/
-
-    cleanup:
-    dbmClient?.close()
-  }
 }
 
 class MongoJava31ClientV0Test extends MongoJava31ClientTest {