1. Use assertThat instead of waitAndAssertTraces

2. Close resources using AutoCleanupExtension
3. Reconstruct the mongoSpan method

Author: xiepuhuan

instrumentation/mongo/mongo-3.1/testing/src/main/java/io/opentelemetry/instrumentation/mongo/v3_1/AbstractMongo31ClientTest.java

@@ -6,6 +6,7 @@
 package io.opentelemetry.instrumentation.mongo.v3_1;
 
 import static java.util.Arrays.asList;
+import static org.assertj.core.api.Assertions.assertThat;
 import static org.assertj.core.api.Assertions.assertThatExceptionOfType;
 
 import com.mongodb.MongoClient;
@@ -219,6 +220,6 @@ void testClientFailure() {
               db.createCollection(createCollectionName());
             });
     // Unfortunately not caught by our instrumentation.
-    testing().waitAndAssertTraces();
+    assertThat(testing().spans().size()).isEqualTo(0);
   }
 }

instrumentation/mongo/mongo-3.7/javaagent/src/test/java/io/opentelemetry/javaagent/instrumentation/mongo/v3_7/MongoClientTest.java

@@ -8,6 +8,7 @@
 import static io.opentelemetry.instrumentation.test.utils.PortUtils.UNUSABLE_PORT;
 import static java.util.Arrays.asList;
 import static java.util.Collections.singletonList;
+import static org.assertj.core.api.Assertions.assertThat;
 import static org.assertj.core.api.Assertions.assertThatThrownBy;
 
 import com.mongodb.MongoClientSettings;
@@ -240,6 +241,6 @@ void testClientFailure() {
             })
         .isInstanceOf(MongoTimeoutException.class);
     // Unfortunately not caught by our instrumentation.
-    testing().waitAndAssertTraces();
+    assertThat(testing().spans().size()).isEqualTo(0);
   }
 }

instrumentation/mongo/mongo-4.0/javaagent/src/test/java/io/opentelemetry/javaagent/instrumentation/mongo/v4_0/Mongo4ReactiveClientTest.java

@@ -16,10 +16,9 @@
 import com.mongodb.reactivestreams.client.MongoCollection;
 import com.mongodb.reactivestreams.client.MongoDatabase;
 import io.opentelemetry.instrumentation.mongo.testing.AbstractMongoClientTest;
+import io.opentelemetry.instrumentation.testing.internal.AutoCleanupExtension;
 import io.opentelemetry.instrumentation.testing.junit.AgentInstrumentationExtension;
 import io.opentelemetry.instrumentation.testing.junit.InstrumentationExtension;
-import java.util.ArrayList;
-import java.util.List;
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.CountDownLatch;
 import java.util.concurrent.ExecutionException;
@@ -41,23 +40,20 @@ class Mongo4ReactiveClientTest extends AbstractMongoClientTest<MongoCollection<D
   @RegisterExtension
   static final InstrumentationExtension testing = AgentInstrumentationExtension.create();
 
-  private MongoClient client;
+  @RegisterExtension static final AutoCleanupExtension cleanup = AutoCleanupExtension.create();
 
-  private final List<AutoCloseable> cleanup = new ArrayList<>();
+  private MongoClient client;
 
   @BeforeAll
   void setup() {
     client = MongoClients.create("mongodb://" + host + ":" + port);
   }
 
   @AfterAll
-  void cleanup() throws Exception {
+  void cleanup() {
     if (client != null) {
       client.close();
     }
-    for (AutoCloseable resource : cleanup) {
-      resource.close();
-    }
   }
 
   @Override
@@ -82,7 +78,7 @@ public void createCollection(String dbName, String collectionName) {
   public void createCollectionNoDescription(String dbName, String collectionName)
       throws InterruptedException {
     MongoClient tmpClient = MongoClients.create("mongodb://" + host + ":" + port);
-    cleanup.add(tmpClient);
+    cleanup.deferCleanup(tmpClient);
     MongoDatabase db = tmpClient.getDatabase(dbName);
     CountDownLatch latch = new CountDownLatch(1);
     db.createCollection(collectionName).subscribe(toSubscriber(o -> latch.countDown()));
@@ -103,7 +99,7 @@ public void createCollectionCallingBuildTwice(String dbName, String collectionNa
                 builder -> builder.hosts(singletonList(new ServerAddress(host, port))));
     settings.build();
     MongoClient tmpClient = MongoClients.create(settings.build());
-    cleanup.add(tmpClient);
+    cleanup.deferCleanup(tmpClient);
     MongoDatabase db = tmpClient.getDatabase(dbName);
     CountDownLatch latch = new CountDownLatch(1);
     db.createCollection(collectionName).subscribe(toSubscriber(o -> latch.countDown()));

instrumentation/mongo/mongo-4.0/javaagent/src/test/java/io/opentelemetry/javaagent/instrumentation/mongo/v4_0/MongoClientTest.java

@@ -71,12 +71,12 @@ public void createCollectionWithAlreadyBuiltClientOptions(String dbName, String
 
   @Override
   public void createCollectionCallingBuildTwice(String dbName, String collectionName) {
-    MongoClientSettings settings =
+    MongoClientSettings.Builder settings =
         MongoClientSettings.builder()
             .applyToClusterSettings(
-                builder -> builder.hosts(singletonList(new ServerAddress(host, port))))
-            .build();
-    MongoDatabase db = MongoClients.create(settings).getDatabase(dbName);
+                builder -> builder.hosts(singletonList(new ServerAddress(host, port))));
+    settings.build();
+    MongoDatabase db = MongoClients.create(settings.build()).getDatabase(dbName);
     db.createCollection(collectionName);
   }