Add a small test utility that mimics Spock's 'cleanup:' block (#3533)

* Add a small test utility that mimics Spock's 'cleanup:' block

* Update testing-common/src/main/java/io/opentelemetry/instrumentation/testing/junit/AutoCleanupExtension.java

Co-authored-by: Anuraag Agrawal <anuraaga@gmail.com>

* Move to internal

* fix broken import

Co-authored-by: Anuraag Agrawal <anuraaga@gmail.com>

instrumentation/rmi/javaagent/src/test/java/io/opentelemetry/javaagent/instrumentation/rmi/RmiTest.java

@@ -13,6 +13,7 @@ import static org.assertj.core.api.Assertions.entry;
 import io.opentelemetry.api.trace.SpanId;
 import io.opentelemetry.api.trace.SpanKind;
 import io.opentelemetry.instrumentation.test.utils.PortUtils;
+import io.opentelemetry.instrumentation.testing.internal.AutoCleanupExtension;
 import io.opentelemetry.instrumentation.testing.junit.AgentInstrumentationExtension;
 import io.opentelemetry.semconv.trace.attributes.SemanticAttributes;
 import java.rmi.registry.LocateRegistry;
@@ -36,6 +37,8 @@ class RmiTest {
   private static Registry serverRegistry;
   private static Registry clientRegistry;
 
+  @RegisterExtension final AutoCleanupExtension autoCleanup = AutoCleanupExtension.create();
+
   @BeforeAll
   static void setUp() throws Exception {
     registryPort = PortUtils.findOpenPort();
@@ -52,6 +55,7 @@ class RmiTest {
   void clientCallCreatesSpans() throws Exception {
     Server server = new Server();
     serverRegistry.rebind(Server.RMI_ID, server);
+    autoCleanup.deferCleanup(() -> serverRegistry.unbind(Server.RMI_ID));
 
     String response =
         runUnderTrace(
@@ -99,14 +103,13 @@ class RmiTest {
                                                     SemanticAttributes.RPC_SERVICE,
                                                     "rmi.app.Server"),
                                                 entry(SemanticAttributes.RPC_METHOD, "hello")))));
-
-    serverRegistry.unbind(Server.RMI_ID);
   }
 
   @Test
   void serverBuiltinMethods() throws Exception {
     Server server = new Server();
     serverRegistry.rebind(Server.RMI_ID, server);
+    autoCleanup.deferCleanup(() -> serverRegistry.unbind(Server.RMI_ID));
 
     server.equals(new Server());
     server.getRef();
@@ -115,14 +118,13 @@ class RmiTest {
     server.getClass();
 
     assertThat(testing.waitForTraces(0)).isEmpty();
-
-    serverRegistry.unbind(Server.RMI_ID);
   }
 
   @Test
   void serviceThrownException() throws Exception {
     Server server = new Server();
     serverRegistry.rebind(Server.RMI_ID, server);
+    autoCleanup.deferCleanup(() -> serverRegistry.unbind(Server.RMI_ID));
 
     Throwable thrown =
         catchThrowableOfType(
@@ -224,14 +226,13 @@ class RmiTest {
                                                 entry(
                                                     SemanticAttributes.RPC_METHOD,
                                                     "exceptional")))));
-
-    serverRegistry.unbind(Server.RMI_ID);
   }
 
   @Test
   void clientCallUsingLegacyStub() throws Exception {
     ServerLegacy server = new ServerLegacy();
     serverRegistry.rebind(ServerLegacy.RMI_ID, server);
+    autoCleanup.deferCleanup(() -> serverRegistry.unbind(ServerLegacy.RMI_ID));
 
     String response =
         runUnderTrace(
@@ -279,7 +280,5 @@ class RmiTest {
                                                     SemanticAttributes.RPC_SERVICE,
                                                     "rmi.app.ServerLegacy"),
                                                 entry(SemanticAttributes.RPC_METHOD, "hello")))));
-
-    serverRegistry.unbind(ServerLegacy.RMI_ID);
   }
 }

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

@@ -0,0 +1,56 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.instrumentation.testing.internal;
+
+import java.util.ArrayList;
+import java.util.List;
+import java.util.Queue;
+import java.util.concurrent.ConcurrentLinkedQueue;
+import org.junit.jupiter.api.extension.AfterEachCallback;
+import org.junit.jupiter.api.extension.ExtensionContext;
+
+/**
+ * A small utility extension that allows deferring executing cleanup code in {@code @Test} methods
+ * until the test has finished. All cleanup callbacks added during the test will always be executed
+ * after it finishes, no matter the outcome. Inspired by Spock's {@code cleanup:} block.
+ */
+public final class AutoCleanupExtension implements AfterEachCallback {
+  private final Queue<AutoCloseable> thingsToCleanUp = new ConcurrentLinkedQueue<>();
+
+  private AutoCleanupExtension() {}
+
+  public static AutoCleanupExtension create() {
+    return new AutoCleanupExtension();
+  }
+
+  /** Add a {@code cleanupAction} to execute after the test finishes. */
+  public void deferCleanup(AutoCloseable cleanupAction) {
+    thingsToCleanUp.add(cleanupAction);
+  }
+
+  @Override
+  public void afterEach(ExtensionContext extensionContext) throws Exception {
+    List<Exception> exceptions = new ArrayList<>();
+    while (!thingsToCleanUp.isEmpty()) {
+      try {
+        thingsToCleanUp.poll().close();
+      } catch (Exception e) {
+        exceptions.add(e);
+      }
+    }
+
+    switch (exceptions.size()) {
+      case 0:
+        return;
+      case 1:
+        throw exceptions.get(0);
+      default:
+        AssertionError allFailures = new AssertionError("Multiple cleanup errors occurred");
+        exceptions.forEach(allFailures::addSuppressed);
+        throw allFailures;
+    }
+  }
+}

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

@@ -0,0 +1,33 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.instrumentation.testing.junit;
+
+import static org.junit.jupiter.api.Assertions.assertEquals;
+
+import io.opentelemetry.instrumentation.testing.internal.AutoCleanupExtension;
+import java.util.concurrent.atomic.AtomicInteger;
+import org.junit.jupiter.api.AfterAll;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.extension.RegisterExtension;
+
+class AutoCleanupExtensionTest {
+
+  static final AtomicInteger count = new AtomicInteger(0);
+
+  @RegisterExtension final AutoCleanupExtension autoCleanup = AutoCleanupExtension.create();
+
+  @Test
+  void shouldRunCleanupAfterTest() {
+    autoCleanup.deferCleanup(count::incrementAndGet);
+
+    assertEquals(0, count.get());
+  }
+
+  @AfterAll
+  static void verifyCount() {
+    assertEquals(1, count.get());
+  }
+}