Enhancing JUnit with Default Timeouts and Stack Traces

As mentioned in other Attivio blogs, we use many tools in our agile development environment, including Ant and JUnit. We have over 4000 JUnit tests and rely on them to keep our code robust and allow us to quickly deploy quality releases. One obstacle we hit when running our continuous and nightly tests was in determining how to deal with test timeouts and long running tests. Currently JUnit supports setting timeouts on a per-test basis with the @Test(timeout=n) annotation. In addition, the JUnit Ant Task has a parameter, timeout=n, which allows for test timeouts managed outside of JUnit.

Both of these timeout solutions have a serious flaw in that neither gives enough information to help debug what went wrong or was going on when the timeout occurred. After getting occasional timeouts on different tests with no way of understanding why, we decided to enhance JUnit with two new features. We added a capability to set a default timeout for all JUnit tests, and a capability to capture the stack trace of all threads active at the time of the timeout.

Luckily, JUnit 4.5 has the ability to easily override some of its core functionality allowing us to improve the timeout behavior.

New Custom JUnit Classes

To add the ability to have a default timeout for all JUnit tests based on a system parameter, -Djunit.default.timeout=n, we created the class JUnit4EnhancedTimeoutClassRunner as a subclass of org.junit.runners.BlockJUnit4ClassRunner.

package com.attivio.junit.runners;
 
import org.junit.Test;
import org.junit.runners.BlockJUnit4ClassRunner;
import org.junit.runners.model.FrameworkMethod;
import org.junit.runners.model.InitializationError;
import org.junit.runners.model.Statement;

import com.attivio.junit.internal.runners.statements.FailAndStackDumpOnTimeout;

/**
 * JUnit timeout class runner that manages tests with timeouts.
 * If the timeout property is set on a test's @Test annotation, then 
 * that timeout is used for this test, otherwise the system
 * property junit.default.timeout is used.  If no timeout is
 * set on the test and no default timeout is set for all tests, then
 * the test will never timeout. 
 *
 */
public class JUnit4EnhancedTimeoutClassRunner extends BlockJUnit4ClassRunner {
/**
 * Constructor
 * @param klass
 * @throws InitializationError
 */
 public JUnit4EnhancedTimeoutClassRunner(Class<?> klass) throws InitializationError {
 super(klass);
 }

@Override
 protected Statement withPotentialTimeout(FrameworkMethod method, Object test, Statement next) {
   long timeout= getTimeout(method.getAnnotation(Test.class));
   if (timeout == 0) {
      /*
       *  If no timeout is specified on the Test annotation, use a default timeout
       *  if the system property junit.default.timeout is set.
       */
      String defaultTimeout = System.getProperty("junit.default.timeout");
      if (defaultTimeout != null) {
        timeout = Long.parseLong(defaultTimeout);
      }
   }
   return timeout > 0 ? new FailAndStackDumpOnTimeout(next, timeout) : next;
}
/**
 * Gets the timeout set on a @Test annotation.
 * @param annotation of the test
 * @return the timeout or 0 if not timeout is set.
 */
private long getTimeout(Test annotation) {
  if (annotation == null)
    return 0;
  return annotation.timeout();
 }
}

To dump the stack traces on all threads when a timeout occurs, we created a new class FailAndStackDumpOnTimeout based on the org.junit.internal.runners.statements.FailOnTimeout class in JUnit 4.

package com.attivio.junit.internal.runners.statements;
import java.util.Map;
import java.util.concurrent.Callable;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.Future;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.TimeoutException;
import org.junit.runners.model.Statement;
/**
 * Class to cause failures on junit tests that time out.  
 * For timeout failures, the stack traces of all threads
 * is appended to the exception.
 * 
 * Enhanced version of {@link FailOnTimeout}
 */
public class FailAndStackDumpOnTimeout extends Statement {
  private Statement fNext;
  private final long fTimeout;
  /**
   * Constructor 
   * @param next statement
   * @param timeout in milliseconds
   */
 public FailAndStackDumpOnTimeout(Statement next, long timeout) {
   fNext= next;
   fTimeout= timeout;
 }
@Override
 public void evaluate() throws Throwable {
   ExecutorService service= Executors.newSingleThreadExecutor();
   Callable<Object> callable= new Callable<Object>() {
     public Object call() throws Exception {
       try {
         fNext.evaluate();
       } catch (Throwable e) {
         throw new ExecutionException(e);
       }
       return null;
     }
   };
   Future<Object> result= service.submit(callable);
   service.shutdown();
   String stackDump = "";
   try {
     boolean terminated= service.awaitTermination(fTimeout,
        TimeUnit.MILLISECONDS);
     if (!terminated) {
        stackDump = getStackTraces();
        service.shutdownNow();
     }
     result.get(0, TimeUnit.MILLISECONDS); // throws the exception if one occurred during the invocation
 } catch (TimeoutException e) {
     if (stackDump.length() > 0) {
       throw new Exception(String.format(
          "test timed out after %d milliseconds\nAll threads active when test timeout occurred:\n%s", fTimeout, stackDump));
     } else {
       throw new Exception(String.format(
          "test timed out after %d milliseconds", fTimeout));
     }
   } catch (ExecutionException e) {
     throw unwrap(e);
  }
 }
/**
 * Gets all thread stack traces.
 * @return string of all thread stack traces
 */
 private String getStackTraces() {
   StringBuilder sb = new StringBuilder();
   Map<Thread, StackTraceElement[]> stacks = Thread.getAllStackTraces();
   for (Thread t : stacks.keySet()) {
     sb.append(t.toString()).append('\n');
     for (StackTraceElement ste : t.getStackTrace()) {
       sb.append("\tat ").append(ste.toString()).append('\n');
     }
     sb.append('\n');
   }
   return sb.toString();
 }
/**
 * Unwraps the exception
 * @param e Throwable exception
 * @return unwrapped exception
 */
 private Throwable unwrap(Throwable e) {
   if (e instanceof ExecutionException)
     return unwrap(e.getCause());
   return e;
 }
}

To have JUnit use our JUnit4EnhancedTimeoutClassRunner class instead of its default BlockJUnit4ClassRunner, we had to create our own version of org.junit.internal.builders.JUnit4Builder and put our jar before the JUnit jar on the classpath. Our builder code is as follows:

package org.junit.internal.builders;
import org.junit.runner.Runner;
import org.junit.runners.model.RunnerBuilder;
import com.attivio.junit.runners.JUnit4EnhancedTimeoutClassRunner;

/**
 * JUnit4Builder override class to create a runner
 * of type {@link JUnit4EnhancedTimeoutClassRunner} by default.
 * This must come first on the classpath before JUnit 4's jar so it
 * is instantiated instead of the default JUnit 4 builder.
 */
public class JUnit4Builder extends RunnerBuilder {
  @Override
   public Runner runnerForClass(Class<?> testClass) throws Throwable {
     // Use the new enhanced timeout class runner
     return new JUnit4EnhancedTimeoutClassRunner(testClass);
  }
}

Finally, to have Ant JUnit use our builder instead of the default JUnit 4 builder, we put our junit-override.jar containing our new JUnit code on the classpath before the junit-4.5.jar. We modified our ant file as follows:

<path id="test.path">
    <path location="build/classes" />
    <path location="build/test" />
    <path location="lib/custom/junit-override.jar" />
    <path location="lib/junit/junit-4.5.jar" />
 </path>

Conclusion

Once we finished enhancing JUnit, we set the default timeout for all of our tests to be 5 minutes by specifying the junit.default.timeout property on the JUnit ant target as follows:

<property name="junit.default.timeout" value="300000" /> <!-- 5 minute default timeout is 300000 -->
<junit printsummary="yes" haltonfailure="no" dir="${basedir}" fork="yes" forkmode="@{forkmode}" errorproperty="tests.failed"
        showoutput="${test.showOutput}" outputtoformatters="${test.outputToFormatters}" tempdir="build/test">
   <sysproperty key="junit.default.timeout" value="${junit.default.timeout}" />
   ...
</junit>

For any tests we expect to take longer than the default of 5 minutes, we specify the expected timeout via the @Test annotation on that test directly.

@Test(timeout=600000)
public void longTest() {
...
}

Author Bio

Since graduating from Carnegie Mellon University with a Master's degree in Information Networking, Greg George has had extensive experience developing high performance, large-scale data processing solutions. Greg was an early employee at Ab Initio Software Corporation where he solely designed and developed their database interface engine for all of the major database vendors. After Ab Initio Greg worked at Lumigent Technologies as a technical leader designing auditing solutions for Oracle, DB2 and Sybase. Greg is a Principal Software Engineer working at Attivio.

TrackBack URI for this entry

Comments (4)