diff --git a/instrumentation/rmi/javaagent/src/test/java/io/opentelemetry/javaagent/instrumentation/rmi/RmiTest.java b/instrumentation/rmi/javaagent/src/test/java/io/opentelemetry/javaagent/instrumentation/rmi/RmiTest.java index c12cfb3b1a..d87589f203 100644 --- a/instrumentation/rmi/javaagent/src/test/java/io/opentelemetry/javaagent/instrumentation/rmi/RmiTest.java +++ b/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); } } diff --git a/testing-common/src/main/java/io/opentelemetry/instrumentation/testing/internal/AutoCleanupExtension.java b/testing-common/src/main/java/io/opentelemetry/instrumentation/testing/internal/AutoCleanupExtension.java new file mode 100644 index 0000000000..c827b4845e --- /dev/null +++ b/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 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 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; + } + } +} diff --git a/testing-common/src/test/java/io/opentelemetry/instrumentation/testing/junit/AutoCleanupExtensionTest.java b/testing-common/src/test/java/io/opentelemetry/instrumentation/testing/junit/AutoCleanupExtensionTest.java new file mode 100644 index 0000000000..9e6e578533 --- /dev/null +++ b/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()); + } +}