Writing Your Own JUnit Extensions Using @Rule

I was surprised to see how many hidden and undocumented features Junit has. In the new release (4.7) you have access to a lot of interesting stuff (still in development for some). One that is really interesting is @Rule.

This annotation allows you to annotate a public field in your test class, which is of type MethodRule. This binding will intercept test method calls like an AOP framework would do and redefine the execution, skip it, or do anything else.

In example, suppose you want to run some concurrency test: you may need to execute your test method on 15 threads each starting at the same time, and then wait for all threads to finish. All this plumbing can now be resumed by one annotation:

@Test
@Concurrent(15)
public void myTestMethod() throws InterruptedException {
    System.out.println("Thread " + Thread.currentThread().getName() + " started !");
    int n = new Random().nextInt(5000);
    System.out.println("Thread " + Thread.currentThread().getName() + " wait " + n + "ms");
    Thread.sleep(n);
    System.out.println("Thread " + Thread.currentThread().getName() + " finished");

But who manages this annotation ? It will be our ManageRule, and it is used like this:

public final class ConcurrentTest {

    @Rule
    public ConcurrentRule concurrentRule = new ConcurrentRule();

    @Test
    @Concurrent(15)
    public void myTestMethod() throws InterruptedException {
        System.out.println("Thread " + Thread.currentThread().getName() + " started !");
        int n = new Random().nextInt(5000);
        System.out.println("Thread " + Thread.currentThread().getName() + " wait " + n + "ms");
        Thread.sleep(n);
        System.out.println("Thread " + Thread.currentThread().getName() + " finished");
    }
}

ManageRule are simplier instances that will intercept all test calls. They are mere instance variable that can benefit of parametrization. In example, we may want in a test class to have a ConcurrentRule with a timeout of 10 seconds, and 2 seconds in another class.

In our case, the ConcurrentRule I made simply check if the test method intercepted has a Concurrent annotation. If yes, it will spawn the number of thread requested and launch the test method for each thread.Here is the annotation:

@Retention(RetentionPolicy.RUNTIME)
@Target({ElementType.METHOD})
public @interface Concurrent {
    int value() default 10;
}

And the MethodRule implementation:

import org.junit.rules.MethodRule;
import org.junit.runners.model.FrameworkMethod;
import org.junit.runners.model.Statement;
import java.util.concurrent.CountDownLatch;

public final class ConcurrentRule implements MethodRule {
    @Override
    public Statement apply(Statement statement, final FrameworkMethod frameworkMethod, final Object o) {
        return new Statement() {
            @Override
            public void evaluate() throws Throwable {
                Concurrent concurrent = frameworkMethod.getAnnotation(Concurrent.class);
                if (concurrent == null)
                    frameworkMethod.invokeExplosively(o);
                else {
                    final String name = frameworkMethod.getName();
                    final Thread[] threads = new Thread[concurrent.value()];
                    final CountDownLatch go = new CountDownLatch(1);
                    final CountDownLatch finished = new CountDownLatch(threads.length);
                    for (int i = 0; i < threads.length; i++) {
                        threads[i] = new Thread(new Runnable() {
                            @Override
                            public void run() {
                                try {
                                    go.await();
                                } catch (InterruptedException e) {
                                    Thread.currentThread().interrupt();
                                }
                                try {
                                    frameworkMethod.invokeExplosively(o);
                                } catch (Throwable throwable) {
                                    if (throwable instanceof RuntimeException)
                                        throw (RuntimeException) throwable;
                                    if (throwable instanceof Error)
                                        throw (Error) throwable;
                                    RuntimeException r = new RuntimeException(throwable.getMessage(), throwable);
                                    r.setStackTrace(throwable.getStackTrace());
                                    throw r;
                                } finally {
                                    finished.countDown();
                                }
                            }
                        }, name + "-Thread-" + i);
                        threads[i].start();
                    }
                    go.countDown();
                    finished.await();
                }
            }
        };
    }
}

From http://blog.mycila.com