Wait for @AfterEach method's VertxTestContext in case a @test method failed

Author: alsin

src/main/java/io/vertx/junit5/VertxExtension.java

@@ -12,10 +12,17 @@
 package io.vertx.junit5;
 
 import io.vertx.core.Vertx;
+import org.junit.jupiter.api.AfterEach;
 import org.junit.jupiter.api.Nested;
-import org.junit.jupiter.api.extension.*;
+import org.junit.jupiter.api.extension.DynamicTestInvocationContext;
+import org.junit.jupiter.api.extension.ExtensionContext;
 import org.junit.jupiter.api.extension.ExtensionContext.Namespace;
 import org.junit.jupiter.api.extension.ExtensionContext.Store;
+import org.junit.jupiter.api.extension.InvocationInterceptor;
+import org.junit.jupiter.api.extension.ParameterContext;
+import org.junit.jupiter.api.extension.ParameterResolutionException;
+import org.junit.jupiter.api.extension.ParameterResolver;
+import org.junit.jupiter.api.extension.ReflectiveInvocationContext;
 
 import java.lang.reflect.Method;
 import java.util.ArrayList;
@@ -164,12 +171,23 @@ public void interceptDynamicTest(Invocation<Void> invocation, DynamicTestInvocat
   @Override
   public void interceptAfterEachMethod(Invocation<Void> invocation, ReflectiveInvocationContext<Method> invocationContext, ExtensionContext extensionContext) throws Throwable {
     invocation.proceed();
-    joinActiveTestContexts(extensionContext);
+    joinActiveTestContexts(invocationContext, extensionContext);
   }
 
   private void joinActiveTestContexts(ExtensionContext extensionContext) throws Exception {
+    joinActiveTestContexts(null, extensionContext);
+  }
+
+  private void joinActiveTestContexts(ReflectiveInvocationContext<Method> invocationContext, ExtensionContext extensionContext) throws Exception {
     if (extensionContext.getExecutionException().isPresent()) {
-      return;
+      final boolean isNotInAfterEachMethod = Optional.ofNullable(invocationContext)
+        .map(ReflectiveInvocationContext::getExecutable)
+        .map(executable -> executable.getAnnotation(AfterEach.class))
+        .isEmpty();
+
+      if (isNotInAfterEachMethod) {
+        return;
+      }
     }
 
     ContextList currentContexts = store(extensionContext).remove(TEST_CONTEXT_KEY, ContextList.class);

src/test/java/io/vertx/junit5/tests/WaitForContextInAfterEachMethodTest.java

@@ -0,0 +1,40 @@
+package io.vertx.junit5.tests;
+
+import io.vertx.core.Vertx;
+import io.vertx.junit5.VertxExtension;
+import io.vertx.junit5.VertxTestContext;
+import org.junit.jupiter.api.AfterAll;
+import org.junit.jupiter.api.AfterEach;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.extension.ExtendWith;
+
+import java.util.concurrent.atomic.AtomicBoolean;
+
+@ExtendWith({VertxExtension.class})
+public class WaitForContextInAfterEachMethodTest {
+    static final AtomicBoolean afterTestContextAwaited = new AtomicBoolean(false);
+
+    @Test
+    void test(final VertxTestContext context) {
+        context.failNow(new RuntimeException("Failing test through context"));
+    }
+
+    @AfterAll
+    static void afterAll(final VertxTestContext context) {
+      if (afterTestContextAwaited.get()) {
+        context.completeNow();
+      } else {
+        context.failNow(new RuntimeException("afterTest context was not awaited"));
+      }
+    }
+
+    @AfterEach
+    void afterTest(final Vertx vertx, final VertxTestContext context) {
+        vertx.setTimer(100, id -> {
+            afterTestContextAwaited.set(true);
+            context.completeNow();
+        });
+
+    }
+
+}