Write output from JUnit ConsoleRunner when process is terminated

src/java/org/pantsbuild/tools/junit/impl/AntJunitXmlReportListener.java

@@ -13,8 +13,6 @@
 import java.io.FilterWriter;
 import java.io.IOException;
 import java.io.Writer;
-import java.net.InetAddress;
-import java.net.UnknownHostException;
 import java.nio.file.Files;
 import java.text.SimpleDateFormat;
 import java.util.Date;
@@ -305,9 +303,9 @@ public void started() {
     }
 
     public void finished() {
-      tests++;
       if (startNs != 0) {
         time = convertTimeSpanNs(System.nanoTime() - startNs);
+        tests++;
       } else {
         time = "0";
       }
@@ -426,6 +424,7 @@ public void testFailure(Failure failure) throws java.lang.Exception {
         if (suite != null) {
           suite.testCases.clear();
           suite.testCases.add(testCase);
+          suite.tests = 1;
           suite.finished();
         }
       } else {
@@ -455,6 +454,8 @@ public void testIgnored(Description description) throws java.lang.Exception {
       suite.incrementSkipped();
       // Ignored tests don't have testStarted and testFinished callbacks so call finish here.
       suite.finished();
+      // finished() counts only started tests, so count ignored tests here so the totals are correct
+      if (!suite.wasStarted()) suite.tests++;
     }
 
     TestCase testCase = getTestCaseFor(description);
@@ -489,6 +490,7 @@ public void testFinished(Description description) throws java.lang.Exception {
   public void testRunFinished(Result result) throws java.lang.Exception {
     for (TestSuite suite : suites.values()) {
       if (suite.wasStarted() && suite.testClass != null) {
+        streamSource.close(suite.testClass); // may not be closed if we're abnormally terminating
         suite.setOut(new String(streamSource.readOut(suite.testClass), Charsets.UTF_8));
         suite.setErr(new String(streamSource.readErr(suite.testClass), Charsets.UTF_8));
       }

src/java/org/pantsbuild/tools/junit/impl/ConsoleRunnerImpl.java

@@ -28,6 +28,8 @@
 import java.util.List;
 import java.util.Map;
 import java.util.Set;
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.TimeUnit;
 import org.apache.commons.io.output.TeeOutputStream;
 import org.junit.runner.Computer;
 import org.junit.runner.Description;
@@ -315,6 +317,11 @@ public byte[] readOut(Class<?> testClass) throws IOException {
     public byte[] readErr(Class<?> testClass) throws IOException {
       return suiteCaptures.get(testClass).readErr();
     }
+
+    @Override
+    public void close(Class<?> testClass) throws IOException {
+      suiteCaptures.get(testClass).close();
+    }
   }
 
   /**
@@ -375,6 +382,8 @@ enum OutputMode {
   private final boolean useExperimentalRunner;
   private final SwappableStream<PrintStream> swappableOut;
   private final SwappableStream<PrintStream> swappableErr;
+  private final Set<Thread> shutdownHooks = Sets.newHashSet(); // for use in tests
+  private Collection<String> testsToRun;
 
   ConsoleRunnerImpl(
       boolean failFast,
@@ -411,12 +420,26 @@ enum OutputMode {
     this.useExperimentalRunner = useExperimentalRunner;
   }
 
+  // by holding on to the tests to run, we can create a runner that is ready to go
+  // without beginning the run (which is useful in the tests for the runner itself)
+  private void setTestsToRun(Collection<String> tests) {
+    this.testsToRun = tests;
+  }
+
   void run(Collection<String> tests) {
+    setTestsToRun(tests);
+    run();
+  }
+
+  @VisibleForTesting
+  public void run() {
     System.setOut(new PrintStream(swappableOut));
     System.setErr(new PrintStream(swappableErr));
 
     JUnitCore core = new JUnitCore();
 
+    int numShutdownHooks = 0;
+
     if (testListener != null) {
       core.addListener(testListener);
     }
@@ -429,8 +452,11 @@ void run(Collection<String> tests) {
         new StreamCapturingListener(outdir, outputMode, swappableOut, swappableErr);
     core.addListener(streamCapturingListener);
 
+    RunListener xmlListener = null;
     if (xmlReport) {
-      core.addListener(new AntJunitXmlReportListener(outdir, streamCapturingListener));
+      xmlListener = new AntJunitXmlReportListener(outdir, streamCapturingListener);
+      core.addListener(xmlListener);
+      numShutdownHooks++; // later we will register a shutdown hook for writing XML output
     }
 
     if (perTestTimer) {
@@ -439,17 +465,38 @@ void run(Collection<String> tests) {
       core.addListener(new ConsoleListener(swappableOut.getOriginal()));
     }
 
-    ShutdownListener shutdownListener = new ShutdownListener(swappableOut.getOriginal());
-    core.addListener(shutdownListener);
+    ShutdownListener consoleShutdownListener =
+      new ShutdownListener(new ConsoleListener(swappableOut.getOriginal()));
+    core.addListener(consoleShutdownListener);
+
     // Wrap test execution with registration of a shutdown hook that will ensure we
     // never exit silently if the VM does.
-    final Thread unexpectedExitHook =
-        createUnexpectedExitHook(shutdownListener, swappableOut.getOriginal());
-    Runtime.getRuntime().addShutdownHook(unexpectedExitHook);
+    numShutdownHooks++;
+    final CountDownLatch haltAfterUnexpectedShutdown = new CountDownLatch(numShutdownHooks);
+    final Thread unexpectedExitHook = createUnexpectedExitHook(
+      consoleShutdownListener,
+      swappableOut.getOriginal(),
+      haltAfterUnexpectedShutdown);
+    addShutdownHook(unexpectedExitHook);
+
+    // handle writing XML output when the tests time out and are terminated by pants
+    if (xmlListener != null) {
+      final ShutdownListener xmlShutdownListener = new ShutdownListener(xmlListener);
+      core.addListener(xmlShutdownListener);
+
+      Thread xmlShutdownHook = new Thread() {
+        @Override
+        public void run() {
+          xmlShutdownListener.unexpectedShutdown();
+          haltAfterUnexpectedShutdown.countDown();
+        }
+      };
+      addShutdownHook(xmlShutdownHook);
+    }
 
     int failures = 1;
     try {
-      Collection<Spec> parsedTests = new SpecParser(tests).parse();
+      Collection<Spec> parsedTests = new SpecParser(testsToRun).parse();
       if (useExperimentalRunner) {
         failures = runExperimental(parsedTests, core);
       } else {
@@ -462,15 +509,20 @@ void run(Collection<String> tests) {
     } finally {
       // If we're exiting via a thrown exception, we'll get a better message by letting it
       // propagate than by halt()ing.
-      Runtime.getRuntime().removeShutdownHook(unexpectedExitHook);
+      removeShutdownHook(unexpectedExitHook);
     }
     exit(failures == 0 ? 0 : 1);
   }
 
   /**
-   * Returns a thread that records a system exit to the listener, and then halts(1).
+   * Returns a thread that records a system exit to the listener,
+   * and then halts(1) after other unexpected exit hooks that use the given CountDownLatch.
    */
-  private Thread createUnexpectedExitHook(final ShutdownListener listener, final PrintStream out) {
+  private Thread createUnexpectedExitHook(
+    final ShutdownListener listener,
+    final PrintStream out,
+    final CountDownLatch haltAfter
+  ) {
     return new Thread() {
       @Override public void run() {
         try {
@@ -480,15 +532,39 @@ private Thread createUnexpectedExitHook(final ShutdownListener listener, final P
           out.println(e);
           e.printStackTrace(out);
         }
-        // This error might be a call to `System.exit(0)` in a test, which we definitely do
-        // not want to go unnoticed.
-        out.println("FATAL: VM exiting unexpectedly.");
-        out.flush();
-        Runtime.getRuntime().halt(1);
+        try {
+          // Other shutdown hooks might still need to write test results,
+          // so wait for them (up to 15 seconds) to finish before halting.
+          haltAfter.countDown();
+          long awaiting = haltAfter.getCount();
+          if (awaiting > 0) out.println("Waiting for " + awaiting + "shutdown hooks to complete");
+          haltAfter.await(15, TimeUnit.SECONDS);
+        } catch (InterruptedException e) {
+          out.println(e);
+          e.printStackTrace(out);
+        } finally {
+          if (callSystemExitOnFinish) {
+            // This error might be a call to `System.exit(0)` in a test, which we definitely do
+            // not want to go unnoticed.
+            out.println("FATAL: VM exiting unexpectedly.");
+            out.flush();
+            Runtime.getRuntime().halt(1);
+          }
+        }
       }
     };
   }
 
+  private void addShutdownHook(Thread hook) {
+    shutdownHooks.add(hook);
+    Runtime.getRuntime().addShutdownHook(hook);
+  }
+
+  private void removeShutdownHook(Thread hook) {
+    shutdownHooks.remove(hook);
+    Runtime.getRuntime().removeShutdownHook(hook);
+  }
+
   private int runExperimental(Collection<Spec> parsedTests, JUnitCore core)
       throws InitializationError {
     Preconditions.checkNotNull(core);
@@ -695,12 +771,32 @@ public int compare(Description o1, Description o2) {
     return filteredRequests;
   }
 
+  @VisibleForTesting
+  void runShutdownHooks() throws InterruptedException {
+    for(Thread hook: shutdownHooks) {
+      hook.start();
+    }
+    for(Thread hook: shutdownHooks) {
+      hook.join(10000); // wait for all hooks to complete, up to 10 seconds each
+    }
+  }
+
   /**
    * Launcher for JUnitConsoleRunner.
    *
    * @param args options from the command line
    */
   public static void main(String[] args) {
+    mainImpl(args).run();
+  }
+
+  /**
+   * As main, but returns the ConsoleRunnerImpl instance and doesn't begin the test run.
+   * For use in tests.
+   *
+   * @param args options from the command line
+   */
+  public static ConsoleRunnerImpl mainImpl(String[] args) {
     /**
      * Command line option bean.
      */
@@ -846,7 +942,8 @@ public void setNumRetries(int numRetries) {
         tests.add(test);
       }
     }
-    runner.run(tests);
+    runner.setTestsToRun(tests);
+    return runner;
   }
 
   /**

src/java/org/pantsbuild/tools/junit/impl/ShutdownListener.java

@@ -3,71 +3,89 @@
 
 package org.pantsbuild.tools.junit.impl;
 
-import java.io.PrintStream;
+import java.util.concurrent.ConcurrentHashMap;
 import org.junit.runner.Description;
 import org.junit.runner.Result;
 import org.junit.runner.notification.Failure;
 import org.junit.runner.notification.RunListener;
 
 /**
- * A listener that keeps track of the current test state with its own result class so it can print
+ * A listener that keeps track of the current test state with its own result class so it can record
  * the state of tests being run if there is unexpected exit during the tests.
  */
-public class ShutdownListener extends ConsoleListener {
+public class ShutdownListener extends RunListener {
   private final Result result = new Result();
   private final RunListener resultListener = result.createListener();
-  private Description currentTestDescription;
+  // holds running tests: Descriptions are added on testStarted and removed on testFinished
+  private ConcurrentHashMap.KeySetView<Description, Boolean> currentDescriptions =
+    ConcurrentHashMap.newKeySet();
+  private RunListener underlying;
 
-  public ShutdownListener(PrintStream out) {
-    super(out);
+
+  public ShutdownListener(RunListener underlying) {
+    this.underlying = underlying;
   }
 
   public void unexpectedShutdown() {
-    if (currentTestDescription != null) {
-      Failure shutdownFailure = new Failure(currentTestDescription,
-          new UnknownError("Abnormal VM exit - test crashed."));
-      testFailure(shutdownFailure);
+    for(Description description : currentDescriptions) {
+      completeTestWithFailure(description);
     }
 
-    // Log the test summary to the Console
-    super.testRunFinished(result);
+    try {
+      resultListener.testRunFinished(result);
+      underlying.testRunFinished(result);
+    } catch (Exception e) {
+      throw new RuntimeException(e);
+    }
+  }
+
+  private void completeTestWithFailure(Description description) {
+    Failure shutdownFailure = new Failure(description,
+      new UnknownError("Abnormal VM exit - test crashed. The test run may have timed out."));
+
+    try {
+      // Mark this test as completed with a failure (finish its lifecycle)
+      resultListener.testFailure(shutdownFailure);
+      resultListener.testFinished(description);
+      underlying.testFailure(shutdownFailure);
+      underlying.testFinished(description);
+    } catch (Exception ignored){}
   }
 
   @Override
   public void testRunStarted(Description description) throws Exception {
-    this.currentTestDescription = description;
     resultListener.testRunStarted(description);
   }
 
   @Override
-  public void testRunFinished(Result result) {
-    try {
-      resultListener.testRunFinished(result);
-    } catch (Exception e) {
-      throw new RuntimeException(e);
-    }
+  public void testStarted(Description description) throws Exception {
+    currentDescriptions.add(description);
+    resultListener.testStarted(description);
+  }
+
+  @Override
+  public void testAssumptionFailure(Failure failure) {
+    resultListener.testAssumptionFailure(failure);
+  }
+
+  @Override
+  public void testRunFinished(Result result) throws Exception {
+    resultListener.testRunFinished(result);
   }
 
   @Override
   public void testFinished(Description description) throws Exception {
+    currentDescriptions.remove(description);
     resultListener.testFinished(description);
   }
 
   @Override
-  public void testFailure(Failure failure) {
-    try {
-      resultListener.testFailure(failure);
-    } catch (Exception e) {
-      throw new RuntimeException(e);
-    }
+  public void testFailure(Failure failure) throws Exception {
+    resultListener.testFailure(failure);
   }
 
   @Override
-  public void testIgnored(Description description) {
-    try {
-      resultListener.testIgnored(description);
-    } catch (Exception e) {
-      throw new RuntimeException(e);
-    }
+  public void testIgnored(Description description) throws Exception {
+    resultListener.testIgnored(description);
   }
 }