Modifying UnifiedTest runner to support collecting Micrometer spans

Author: nhachicha

driver-core/src/test/resources/specifications

@@ -38,6 +38,8 @@ dependencies {
 
     // lambda testing
     testImplementation(libs.aws.lambda.core)
+
+    testImplementation("io.micrometer:micrometer-tracing-test:1.4.5")
 }
 
 configureMavenPublication {

driver-sync/build.gradle.kts

@@ -66,6 +66,9 @@
 import com.mongodb.lang.NonNull;
 import com.mongodb.lang.Nullable;
 import com.mongodb.logging.TestLoggingInterceptor;
+import com.mongodb.tracing.MicrometerTracer;
+import io.micrometer.tracing.Tracer;
+import io.micrometer.tracing.test.simple.SimpleTracer;
 import org.bson.BsonArray;
 import org.bson.BsonBoolean;
 import org.bson.BsonDocument;
@@ -112,7 +115,7 @@ public final class Entities {
     private static final Set<String> SUPPORTED_CLIENT_ENTITY_OPTIONS = new HashSet<>(
             asList(
                     "id", "uriOptions", "serverApi", "useMultipleMongoses", "storeEventsAsEntities",
-                    "observeEvents", "observeLogMessages", "observeSensitiveCommands", "ignoreCommandMonitoringEvents"));
+                    "observeEvents", "observeLogMessages", "observeSensitiveCommands", "ignoreCommandMonitoringEvents", "tracing"));
     private final Set<String> entityNames = new HashSet<>();
     private final Map<String, ExecutorService> threads = new HashMap<>();
     private final Map<String, ArrayList<Future<?>>> tasks = new HashMap<>();
@@ -126,6 +129,7 @@ public final class Entities {
     private final Map<String, ClientEncryption> clientEncryptions = new HashMap<>();
     private final Map<String, TestCommandListener> clientCommandListeners = new HashMap<>();
     private final Map<String, TestLoggingInterceptor> clientLoggingInterceptors = new HashMap<>();
+    private final Map<String, Tracer> clientTracing = new HashMap<>();
     private final Map<String, TestConnectionPoolListener> clientConnectionPoolListeners = new HashMap<>();
     private final Map<String, TestServerListener> clientServerListeners = new HashMap<>();
     private final Map<String, TestClusterListener> clientClusterListeners = new HashMap<>();
@@ -294,6 +298,10 @@ public TestLoggingInterceptor getClientLoggingInterceptor(final String id) {
         return getEntity(id + "-logging-interceptor", clientLoggingInterceptors, "logging interceptor");
     }
 
+    public Tracer getClientTracer(final String id) {
+        return getEntity(id + "-tracing", clientTracing, "micrometer tracing");
+    }
+
     public TestConnectionPoolListener getConnectionPoolListener(final String id) {
         return getEntity(id + "-connection-pool-listener", clientConnectionPoolListeners, "connection pool listener");
     }
@@ -604,6 +612,14 @@ private void initClient(final BsonDocument entity, final String id,
             }
             clientSettingsBuilder.serverApi(serverApiBuilder.build());
         }
+
+        if (entity.containsKey("tracing") && entity.getBoolean("tracing").getValue()) {
+            Tracer tracer = new SimpleTracer();
+            putEntity(id + "-tracing", tracer, clientTracing);
+
+            clientSettingsBuilder.tracer(new MicrometerTracer(tracer));
+        }
+
         MongoClientSettings clientSettings = clientSettingsBuilder.build();
 
         if (entity.containsKey("observeLogMessages")) {

driver-sync/src/test/functional/com/mongodb/client/unified/Entities.java

@@ -43,6 +43,10 @@
 import com.mongodb.lang.Nullable;
 import com.mongodb.logging.TestLoggingInterceptor;
 import com.mongodb.test.AfterBeforeParameterResolver;
+import io.micrometer.tracing.Tracer;
+import io.micrometer.tracing.test.simple.SimpleSpan;
+import io.micrometer.tracing.test.simple.SimpleTracer;
+import io.micrometer.tracing.test.simple.SpanAssert;
 import org.bson.BsonArray;
 import org.bson.BsonBoolean;
 import org.bson.BsonDocument;
@@ -67,8 +71,10 @@
 import java.util.Arrays;
 import java.util.Collection;
 import java.util.Collections;
+import java.util.Deque;
 import java.util.HashSet;
 import java.util.List;
+import java.util.Map;
 import java.util.Set;
 import java.util.concurrent.ExecutionException;
 import java.util.concurrent.ExecutorService;
@@ -375,6 +381,11 @@ public void shouldPassAllOutcomes(
                 }
                 compareLogMessages(rootContext, definition, tweaks);
             }
+
+            if (definition.containsKey("expectTracingSpans")) {
+                compareTracingSpans(definition);
+            }
+
         } catch (TestAbortedException e) {
             // if a test is ignored, we do not retry
             throw e;
@@ -455,6 +466,42 @@ private void compareLogMessages(final UnifiedTestContext rootContext, final Bson
         }
     }
 
+    private void compareTracingSpans(final BsonDocument definition) {
+        BsonDocument curTracingSpansForClient = definition.getDocument("expectTracingSpans");
+        String clientId = curTracingSpansForClient.getString("client").getValue();
+
+        // Get the tracer for the client
+        Tracer micrometerTracer = entities.getClientTracer(clientId);
+        SimpleTracer simpleTracer = (SimpleTracer) micrometerTracer;
+
+        // Get the list of expected spans
+        BsonArray expectedSpans = curTracingSpansForClient.getArray("spans");
+
+        // Get the actual reported spans
+        Deque<SimpleSpan> reportedSpans = simpleTracer.getSpans();
+
+        // First assert that we have at least the number of expected spans
+        assertEquals(reportedSpans.size(), expectedSpans.size(), "Expected at least " + expectedSpans.size() + " spans, but found " + reportedSpans.size());
+
+        for (BsonValue expectedSpan : expectedSpans) {
+            BsonDocument expectedSpanDoc = expectedSpan.asDocument();
+            String expectedName = expectedSpanDoc.getString("name").getValue();
+            BsonDocument expectedTags = expectedSpanDoc.getDocument("tags");
+
+            SimpleSpan reportedSpan = reportedSpans.pop();
+
+            SpanAssert.assertThat(reportedSpan)
+                    .hasNameEqualTo(expectedName);
+
+            if (!expectedTags.isEmpty()) {
+                for (Map.Entry<String, BsonValue> tag : expectedTags.entrySet()) {
+                    SpanAssert.assertThat(reportedSpan)
+                            .hasTag(tag.getKey(), tag.getValue().asString().getValue());
+                }
+            }
+        }
+    }
+
     private void assertOutcome(final UnifiedTestContext context) {
         for (BsonValue cur : definition.getArray("outcome")) {
             BsonDocument curDocument = cur.asDocument();