Add AutoCleanupExtension deferAfterAll

Author: trask

.github/agents/knowledge/testing-general-patterns.md

@@ -17,14 +17,19 @@
 - In JUnit tests, when an `AutoCloseable` is intended to remain live for most or all of the test
   and only needs cleanup at test end, prefer `AutoCleanupExtension` with `deferCleanup(...)`
   over wrapping most of the test body in try-with-resources.
+- For resources created in `@BeforeAll` or other class-scoped setup, prefer
+  `AutoCleanupExtension` with `deferAfterAll(...)` over nested `@AfterAll` cleanup
+  chains.
 - Reuse an existing `cleanup` extension when one is already in scope.
   Otherwise, add a `@RegisterExtension` field when the deferred-cleanup pattern improves
   clarity or avoids wrapping most of the test body.
 - Keep try-with-resources for semantically scoped resources whose lifetime must end before the
   rest of the test continues, such as `Scope` / `Context.makeCurrent()`, Mockito
   `MockedStatic`, and short-lived readers, writers, streams, response bodies, or parsers.
-- Do not use `AutoCleanupExtension` in non-JUnit helper methods, `@BeforeAll`, or other shared
-  setup code where cleanup is not tied to a single test method.
+- Keep `AutoCleanupExtension` usage scoped to JUnit-managed test classes; do not introduce it in
+  generic helper utilities or other non-JUnit code.
+- In `@BeforeAll` or other shared setup code where cleanup is not tied to a single test method,
+  use `deferAfterAll(...)` instead of `deferCleanup(...)`.
 - If the test intentionally closes the resource mid-test or asserts behavior around explicit
   close, keep the direct close or try-with-resources in the test body.
 

instrumentation/kafka/kafka-clients/kafka-clients-0.11/testing/src/main/java/io/opentelemetry/instrumentation/kafkaclients/common/v0_11/internal/KafkaClientBaseTest.java

@@ -28,6 +28,7 @@
 import static org.assertj.core.api.Assertions.assertThat;
 
 import io.opentelemetry.api.common.AttributeKey;
+import io.opentelemetry.instrumentation.testing.internal.AutoCleanupExtension;
 import io.opentelemetry.sdk.testing.assertj.AttributeAssertion;
 import java.time.Duration;
 import java.util.ArrayList;
@@ -53,9 +54,9 @@
 import org.apache.kafka.common.serialization.StringSerializer;
 import org.assertj.core.api.AbstractLongAssert;
 import org.assertj.core.api.AbstractStringAssert;
-import org.junit.jupiter.api.AfterAll;
 import org.junit.jupiter.api.BeforeAll;
 import org.junit.jupiter.api.TestInstance;
+import org.junit.jupiter.api.extension.RegisterExtension;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 import org.testcontainers.containers.output.Slf4jLogConsumer;
@@ -67,6 +68,8 @@
 public abstract class KafkaClientBaseTest {
   private static final Logger logger = LoggerFactory.getLogger(KafkaClientBaseTest.class);
 
+  @RegisterExtension final AutoCleanupExtension cleanup = AutoCleanupExtension.create();
+
   protected static final String SHARED_TOPIC = "shared.topic";
   protected static final AttributeKey<String> MESSAGING_CLIENT_ID =
       AttributeKey.stringKey("messaging.client_id");
@@ -84,13 +87,15 @@ public abstract class KafkaClientBaseTest {
 
   @BeforeAll
   void setupClass() throws ExecutionException, InterruptedException, TimeoutException {
-    kafka =
+    KafkaContainer kafkaContainer =
         new KafkaContainer(DockerImageName.parse("apache/kafka:3.8.0"))
             .withEnv("KAFKA_HEAP_OPTS", "-Xmx256m")
             .withLogConsumer(new Slf4jLogConsumer(logger))
             .waitingFor(Wait.forLogMessage(".*started \\(kafka.server.Kafka.*Server\\).*", 1))
             .withStartupTimeout(Duration.ofMinutes(1));
-    kafka.start();
+    kafka = kafkaContainer;
+    cleanup.deferAfterAll(kafkaContainer::stop);
+    kafkaContainer.start();
 
     // create test topic
     HashMap<String, Object> adminProps = new HashMap<>();
@@ -104,8 +109,10 @@ void setupClass() throws ExecutionException, InterruptedException, TimeoutExcept
     }
 
     producer = new KafkaProducer<>(producerProps());
+    cleanup.deferAfterAll(producer);
 
     consumer = new KafkaConsumer<>(consumerProps());
+    cleanup.deferAfterAll(consumer);
 
     consumer.subscribe(
         singletonList(SHARED_TOPIC),
@@ -144,17 +151,6 @@ public Map<String, Object> producerProps() {
     return props;
   }
 
-  @AfterAll
-  void cleanupClass() {
-    if (producer != null) {
-      producer.close();
-    }
-    if (consumer != null) {
-      consumer.close();
-    }
-    kafka.stop();
-  }
-
   public void awaitUntilConsumerIsReady() throws InterruptedException {
     if (consumerReady.await(0, SECONDS)) {
       return;

instrumentation/opensearch/opensearch-rest-common/testing/src/main/java/io/opentelemetry/javaagent/instrumentation/opensearch/rest/AbstractOpenSearchRestTest.java

@@ -24,15 +24,16 @@
 import static org.assertj.core.api.Assertions.assertThat;
 
 import io.opentelemetry.api.trace.SpanKind;
+import io.opentelemetry.instrumentation.testing.internal.AutoCleanupExtension;
 import io.opentelemetry.instrumentation.testing.junit.InstrumentationExtension;
 import java.io.IOException;
 import java.net.URI;
 import java.util.concurrent.CountDownLatch;
 import java.util.concurrent.atomic.AtomicReference;
-import org.junit.jupiter.api.AfterAll;
 import org.junit.jupiter.api.BeforeAll;
 import org.junit.jupiter.api.Test;
 import org.junit.jupiter.api.TestInstance;
+import org.junit.jupiter.api.extension.RegisterExtension;
 import org.opensearch.client.Request;
 import org.opensearch.client.Response;
 import org.opensearch.client.ResponseListener;
@@ -44,6 +45,8 @@
 @TestInstance(TestInstance.Lifecycle.PER_CLASS)
 public abstract class AbstractOpenSearchRestTest {
 
+  @RegisterExtension final AutoCleanupExtension cleanup = AutoCleanupExtension.create();
+
   protected OpensearchContainer opensearch;
   protected RestClient client;
   protected URI httpHost;
@@ -58,26 +61,20 @@ public abstract class AbstractOpenSearchRestTest {
 
   @BeforeAll
   void setUp() throws Exception {
-    opensearch =
+    OpensearchContainer opensearchContainer =
         new OpensearchContainer(DockerImageName.parse("opensearchproject/opensearch:1.3.6"))
             .withSecurityEnabled();
+    opensearch = opensearchContainer;
+    cleanup.deferAfterAll(opensearchContainer::stop);
     // limit memory usage and disable Log4j JMX to avoid cgroup detection issues in containers
-    opensearch.withEnv(
+    opensearchContainer.withEnv(
         "OPENSEARCH_JAVA_OPTS",
         "-Xmx256m -Xms256m -Dlog4j2.disableJmx=true -Dlog4j2.disable.jmx=true -XX:-UseContainerSupport");
-    opensearch.start();
-    httpHost = URI.create(opensearch.getHttpHostAddress());
+    opensearchContainer.start();
+    httpHost = URI.create(opensearchContainer.getHttpHostAddress());
 
     client = buildRestClient();
-  }
-
-  @AfterAll
-  void tearDown() throws IOException {
-    try {
-      client.close();
-    } finally {
-      opensearch.stop();
-    }
+    cleanup.deferAfterAll(client);
   }
 
   @Test

instrumentation/reactor/reactor-kafka-1.0/testing/src/main/java/io/opentelemetry/javaagent/instrumentation/reactor/kafka/v1_0/AbstractReactorKafkaTest.java

@@ -45,7 +45,6 @@
 import org.apache.kafka.common.serialization.StringSerializer;
 import org.assertj.core.api.AbstractLongAssert;
 import org.assertj.core.api.AbstractStringAssert;
-import org.junit.jupiter.api.AfterAll;
 import org.junit.jupiter.api.BeforeAll;
 import org.junit.jupiter.api.extension.RegisterExtension;
 import org.slf4j.Logger;
@@ -80,31 +79,21 @@ public abstract class AbstractReactorKafkaTest {
 
   @BeforeAll
   static void setUpAll() {
-    kafka =
+    KafkaContainer kafkaContainer =
         new KafkaContainer(DockerImageName.parse("apache/kafka:3.8.0"))
             .withEnv("KAFKA_HEAP_OPTS", "-Xmx256m")
             .withLogConsumer(new Slf4jLogConsumer(logger))
             .waitingFor(Wait.forLogMessage(".*started \\(kafka.server.Kafka.*Server\\).*", 1))
             .withStartupTimeout(Duration.ofMinutes(1));
-    kafka.start();
+    kafka = kafkaContainer;
+    cleanup.deferAfterAll(kafkaContainer::stop);
+    kafkaContainer.start();
 
     sender = KafkaSender.create(senderOptions());
+    cleanup.deferAfterAll(sender::close);
     receiver = KafkaReceiver.create(receiverOptions());
   }
 
-  @AfterAll
-  static void tearDownAll() {
-    try {
-      if (sender != null) {
-        sender.close();
-      }
-    } finally {
-      if (kafka != null) {
-        kafka.stop();
-      }
-    }
-  }
-
   @SuppressWarnings("unchecked")
   private static SenderOptions<String, String> senderOptions() {
     Properties props = new Properties();

instrumentation/spring/spring-pulsar-1.0/testing/src/main/java/io/opentelemetry/instrumentation/spring/pulsar/v1_0/AbstractSpringPulsarTest.java

@@ -21,6 +21,7 @@
 import static org.assertj.core.api.Assertions.assertThat;
 
 import io.opentelemetry.instrumentation.testing.GlobalTraceUtil;
+import io.opentelemetry.instrumentation.testing.internal.AutoCleanupExtension;
 import io.opentelemetry.instrumentation.testing.junit.AgentInstrumentationExtension;
 import io.opentelemetry.instrumentation.testing.junit.InstrumentationExtension;
 import io.opentelemetry.sdk.testing.assertj.AttributeAssertion;
@@ -33,7 +34,6 @@
 import org.apache.pulsar.client.api.PulsarClientException;
 import org.assertj.core.api.AbstractLongAssert;
 import org.assertj.core.api.AbstractStringAssert;
-import org.junit.jupiter.api.AfterAll;
 import org.junit.jupiter.api.BeforeAll;
 import org.junit.jupiter.api.Test;
 import org.junit.jupiter.api.extension.RegisterExtension;
@@ -52,6 +52,9 @@ public abstract class AbstractSpringPulsarTest {
   @RegisterExtension
   protected static final InstrumentationExtension testing = AgentInstrumentationExtension.create();
 
+  @RegisterExtension
+  static final AutoCleanupExtension cleanup = AutoCleanupExtension.create();
+
   private static final boolean EXPERIMENTAL_ATTRIBUTES =
       Boolean.getBoolean("otel.instrumentation.pulsar.experimental-span-attributes");
   private static final DockerImageName DEFAULT_IMAGE_NAME =
@@ -74,6 +77,7 @@ static void setUp() throws PulsarClientException {
         new PulsarContainer(DEFAULT_IMAGE_NAME)
             .withEnv("PULSAR_MEM", "-Xmx128m")
             .withStartupTimeout(Duration.ofMinutes(2));
+    cleanup.deferAfterAll(pulsarContainer::stop);
     pulsarContainer.start();
     brokerHost = pulsarContainer.getHost();
     brokerPort = pulsarContainer.getMappedPort(6650);
@@ -85,9 +89,11 @@ static void setUp() throws PulsarClientException {
     props.put("spring.pulsar.consumer.subscription.initial-position", "earliest");
     app.setDefaultProperties(props);
     applicationContext = app.run();
+    cleanup.deferAfterAll(applicationContext);
     pulsarTemplate = applicationContext.getBean(PulsarTemplate.class);
 
     client = PulsarClient.builder().serviceUrl(pulsarContainer.getPulsarBrokerUrl()).build();
+    cleanup.deferAfterAll(client);
   }
 
   @Test
@@ -101,22 +107,6 @@ void testSpringPulsar() throws PulsarClientException, InterruptedException {
     assertSpringPulsar();
   }
 
-  @AfterAll
-  static void teardown() throws PulsarClientException {
-    if (applicationContext != null) {
-      applicationContext.close();
-    }
-    try {
-      if (client != null) {
-        client.close();
-      }
-    } finally {
-      if (pulsarContainer != null) {
-        pulsarContainer.stop();
-      }
-    }
-  }
-
   protected abstract void assertSpringPulsar();
 
   protected List<AttributeAssertion> publishAttributes() {

testing-common/src/main/java/io/opentelemetry/instrumentation/testing/internal/AutoCleanupExtension.java

@@ -6,9 +6,12 @@
 package io.opentelemetry.instrumentation.testing.internal;
 
 import java.util.ArrayList;
+import java.util.Deque;
 import java.util.List;
 import java.util.Queue;
+import java.util.concurrent.ConcurrentLinkedDeque;
 import java.util.concurrent.ConcurrentLinkedQueue;
+import org.junit.jupiter.api.extension.AfterAllCallback;
 import org.junit.jupiter.api.extension.AfterEachCallback;
 import org.junit.jupiter.api.extension.ExtensionContext;
 
@@ -20,8 +23,9 @@
  * <p>This class is internal and is hence not for public use. Its APIs are unstable and can change
  * at any time.
  */
-public class AutoCleanupExtension implements AfterEachCallback {
+public class AutoCleanupExtension implements AfterEachCallback, AfterAllCallback {
   private final Queue<AutoCloseable> thingsToCleanUp = new ConcurrentLinkedQueue<>();
+  private final Deque<AutoCloseable> thingsToCleanUpAfterAll = new ConcurrentLinkedDeque<>();
 
   private AutoCleanupExtension() {}
 
@@ -34,6 +38,11 @@ public void deferCleanup(AutoCloseable cleanupAction) {
     thingsToCleanUp.add(cleanupAction);
   }
 
+  /** Add a {@code cleanupAction} to execute after the test class finishes. */
+  public void deferAfterAll(AutoCloseable cleanupAction) {
+    thingsToCleanUpAfterAll.push(cleanupAction);
+  }
+
   @Override
   public void afterEach(ExtensionContext extensionContext) throws Exception {
     List<Exception> exceptions = new ArrayList<>();
@@ -44,7 +53,23 @@ public void afterEach(ExtensionContext extensionContext) throws Exception {
         exceptions.add(e);
       }
     }
+    throwIfAny(exceptions);
+  }
+
+  @Override
+  public void afterAll(ExtensionContext extensionContext) throws Exception {
+    List<Exception> exceptions = new ArrayList<>();
+    while (!thingsToCleanUpAfterAll.isEmpty()) {
+      try {
+        thingsToCleanUpAfterAll.pop().close();
+      } catch (Exception e) {
+        exceptions.add(e);
+      }
+    }
+    throwIfAny(exceptions);
+  }
 
+  private static void throwIfAny(List<Exception> exceptions) throws Exception {
     switch (exceptions.size()) {
       case 0:
         return;

testing-common/src/test/java/io/opentelemetry/instrumentation/testing/junit/AutoCleanupExtensionTest.java

@@ -30,4 +30,18 @@ void shouldRunCleanupAfterTest() {
   static void verifyCount() {
     assertThat(count.get()).isEqualTo(1);
   }
+
+  @Test
+  void shouldRunCleanupAfterAll() throws Exception {
+    AutoCleanupExtension cleanup = AutoCleanupExtension.create();
+    AtomicInteger countAfterAll = new AtomicInteger(0);
+
+    cleanup.deferAfterAll(countAfterAll::incrementAndGet);
+
+    assertThat(countAfterAll.get()).isEqualTo(0);
+
+    cleanup.afterAll(null);
+
+    assertThat(countAfterAll.get()).isEqualTo(1);
+  }
 }