testing: Have GrpcCleanupRule extend ExternalResource

testing/src/main/java/io/grpc/testing/GrpcCleanupRule.java

@@ -26,20 +26,18 @@
 import io.grpc.ManagedChannel;
 import io.grpc.Server;
 import java.util.ArrayList;
-import java.util.Arrays;
 import java.util.List;
 import java.util.concurrent.TimeUnit;
 import javax.annotation.Nonnull;
 import javax.annotation.concurrent.NotThreadSafe;
-import org.junit.rules.TestRule;
+import org.junit.rules.ExternalResource;
 import org.junit.runner.Description;
-import org.junit.runners.model.MultipleFailureException;
 import org.junit.runners.model.Statement;
 
 /**
- * A JUnit {@link TestRule} that can register gRPC resources and manages its automatic release at
- * the end of the test. If any of the resources registered to the rule can not be successfully
- * released, the test will fail.
+ * A JUnit {@link ExternalResource} that can register gRPC resources and manages its automatic
+ * release at the end of the test. If any of the resources registered to the rule can not be
+ * successfully released, the test will fail.
  *
  * <p>Example usage:
  * <pre>{@code @Rule public final GrpcCleanupRule grpcCleanup = new GrpcCleanupRule();
@@ -73,13 +71,13 @@
  */
 @ExperimentalApi("https://github.com/grpc/grpc-java/issues/2488")
 @NotThreadSafe
-public final class GrpcCleanupRule implements TestRule {
+public final class GrpcCleanupRule extends ExternalResource {
 
   private final List<Resource> resources = new ArrayList<>();
   private long timeoutNanos = TimeUnit.SECONDS.toNanos(10L);
   private Stopwatch stopwatch = Stopwatch.createUnstarted();
 
-  private Throwable firstException;
+  private boolean abruptShutdown;
 
   /**
    * Sets a positive total time limit for the automatic resource cleanup. If any of the resources
@@ -144,71 +142,70 @@ void register(Resource resource) {
     resources.add(resource);
   }
 
+  // The class extends ExternalResource so it can be used in JUnit 5. But JUnit 5 will only call
+  // before() and after(), thus code cannot assume this method will be called.
   @Override
   public Statement apply(final Statement base, Description description) {
-    return new Statement() {
+    return super.apply(new Statement() {
       @Override
       public void evaluate() throws Throwable {
-        firstException = null;
+        abruptShutdown = false;
         try {
           base.evaluate();
         } catch (Throwable t) {
-          firstException = t;
-
-          try {
-            teardown();
-          } catch (Throwable t2) {
-            throw new MultipleFailureException(Arrays.asList(t, t2));
-          }
-
+          abruptShutdown = true;
           throw t;
         }
-
-        teardown();
-        if (firstException != null) {
-          throw firstException;
-        }
       }
-    };
+    }, description);
   }
 
   /**
    * Releases all the registered resources.
    */
-  private void teardown() {
+  @Override
+  protected void after() {
     stopwatch.reset();
     stopwatch.start();
 
-    if (firstException == null) {
+    InterruptedException interrupted = null;
+    if (!abruptShutdown) {
       for (int i = resources.size() - 1; i >= 0; i--) {
         resources.get(i).cleanUp();
       }
+
+      for (int i = resources.size() - 1; i >= 0; i--) {
+        try {
+          boolean released = resources.get(i).awaitReleased(
+              timeoutNanos - stopwatch.elapsed(TimeUnit.NANOSECONDS), TimeUnit.NANOSECONDS);
+          if (released) {
+            resources.remove(i);
+          }
+        } catch (InterruptedException e) {
+          Thread.currentThread().interrupt();
+          interrupted = e;
+          break;
+        }
+      }
     }
 
-    for (int i = resources.size() - 1; i >= 0; i--) {
-      if (firstException != null) {
+    if (!resources.isEmpty()) {
+      for (int i = resources.size() - 1; i >= 0; i--) {
         resources.get(i).forceCleanUp();
-        continue;
       }
 
       try {
-        boolean released = resources.get(i).awaitReleased(
-            timeoutNanos - stopwatch.elapsed(TimeUnit.NANOSECONDS), TimeUnit.NANOSECONDS);
-        if (!released) {
-          firstException = new AssertionError(
-              "Resource " + resources.get(i) + " can not be released in time at the end of test");
+        if (interrupted != null) {
+          throw new AssertionError(
+              "Thread interrupted before resources gracefully released", interrupted);
+        } else if (!abruptShutdown) {
+          throw new AssertionError(
+            "Resources could not be released in time at the end of test: " + resources);
         }
-      } catch (InterruptedException e) {
-        Thread.currentThread().interrupt();
-        firstException = e;
-      }
-
-      if (firstException != null) {
-        resources.get(i).forceCleanUp();
+      } finally {
+        resources.clear();
       }
     }
-
-    resources.clear();
   }
 
   @VisibleForTesting

testing/src/test/java/io/grpc/testing/GrpcCleanupRuleTest.java

@@ -248,13 +248,13 @@ public void multiResource_awaitReleasedFails() throws Throwable {
 
     inOrder.verify(resource3).awaitReleased(anyLong(), any(TimeUnit.class));
     inOrder.verify(resource2).awaitReleased(anyLong(), any(TimeUnit.class));
+    inOrder.verify(resource1).awaitReleased(anyLong(), any(TimeUnit.class));
     inOrder.verify(resource2).forceCleanUp();
-    inOrder.verify(resource1).forceCleanUp();
 
     inOrder.verifyNoMoreInteractions();
 
     verify(resource3, never()).forceCleanUp();
-    verify(resource1, never()).awaitReleased(anyLong(), any(TimeUnit.class));
+    verify(resource1, never()).forceCleanUp();
   }
 
   @Test
@@ -280,7 +280,7 @@ public void multiResource_awaitReleasedInterrupted() throws Throwable {
     boolean cleanupFailed = false;
     try {
       grpcCleanup.apply(statement, null /* description*/).evaluate();
-    } catch (InterruptedException e) {
+    } catch (Throwable e) {
       cleanupFailed = true;
     }
 
@@ -381,15 +381,16 @@ public void multiResource_timeoutCalculation_customTimeout() throws Throwable {
   @Test
   public void baseTestFailsThenCleanupFails() throws Throwable {
     // setup
-    Exception baseTestFailure = new Exception();
+    Exception baseTestFailure = new Exception("base test failure");
+    Exception cleanupFailure = new RuntimeException("force cleanup failed");
 
     Statement statement = mock(Statement.class);
     doThrow(baseTestFailure).when(statement).evaluate();
 
     Resource resource1 = mock(Resource.class);
     Resource resource2 = mock(Resource.class);
     Resource resource3 = mock(Resource.class);
-    doThrow(new RuntimeException()).when(resource2).forceCleanUp();
+    doThrow(cleanupFailure).when(resource2).forceCleanUp();
 
     InOrder inOrder = inOrder(statement, resource1, resource2, resource3);
     GrpcCleanupRule grpcCleanup = new GrpcCleanupRule();
@@ -407,8 +408,14 @@ public void baseTestFailsThenCleanupFails() throws Throwable {
     }
 
     // verify
-    assertThat(failure).isInstanceOf(MultipleFailureException.class);
-    assertSame(baseTestFailure, ((MultipleFailureException) failure).getFailures().get(0));
+    if (failure instanceof MultipleFailureException) {
+      // JUnit 4.13+
+      assertThat(((MultipleFailureException) failure).getFailures())
+          .containsExactly(baseTestFailure, cleanupFailure);
+    } else {
+      // JUnit 4.12. Suffers from https://github.com/junit-team/junit4/issues/1334
+      assertThat(failure).isSameInstanceAs(cleanupFailure);
+    }
 
     inOrder.verify(statement).evaluate();
     inOrder.verify(resource3).forceCleanUp();