PatternMatchingSuite.java

package org.doxla.junit4;

import org.junit.runner.Runner;
import org.junit.runner.notification.RunNotifier;
import org.junit.runners.Suite;
import org.junit.runners.model.RunnerBuilder;
import org.springframework.core.type.classreading.CachingMetadataReaderFactory;
import org.springframework.core.type.classreading.MetadataReaderFactory;

import java.lang.annotation.ElementType;
import java.lang.annotation.Retention;
import java.lang.annotation.RetentionPolicy;
import java.lang.annotation.Target;
import java.util.ArrayList;
import java.util.List;
import java.util.Set;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.TimeUnit;

import static java.util.Arrays.asList;

public class PatternMatchingSuite extends Suite {

    private static final MetadataReaderFactory metadataReaderFactory = new CachingMetadataReaderFactory();

    private static Class<?>[] getClassesFromAnnotation(Class<?> klass) {
        SuiteClassPattern annotation = getSuiteClassAnnotation(klass);
        Set<Class<?>> classes = PatternMatchingClassCollector.getClassesFromPattern(annotation.include());

        classes.removeAll(asList(annotation.exclude()));

        return classes.toArray(new Class<?>[]{});
    }

    private static SuiteClassPattern getSuiteClassAnnotation(Class<?> klass) {
        SuiteClassPattern annotation = klass.getAnnotation(SuiteClassPattern.class);
        if (annotation == null)
            throw new RuntimeException(String.format("class '%s' must have a SuiteClassPattern annotation", klass.getName()));
        return annotation;
    }


    @Retention(RetentionPolicy.RUNTIME)
    @Target(ElementType.TYPE)
    public @interface SuiteClassPattern {
        public String include();

        public Class<?>[] exclude() default {};

        public boolean parallel() default false;
    }

    private final List<Runner> runners = new ArrayList<Runner>();

    public PatternMatchingSuite(Class<?> testClass, RunnerBuilder builder) throws Throwable {
        super(testClass, builder);
        for (Class<?> childTest : getClassesFromAnnotation(testClass)) {
            runners.add(builder.runnerForClass(childTest));
        }
    }

    public boolean shouldRunInParallel(Class<?> clazz) {
        return getSuiteClassAnnotation(clazz).parallel();
    }

    @Override
    public void run(final RunNotifier notifier) {
        if (shouldRunInParallel(getTestClass().getJavaClass())) {
            runTestsInParallel(notifier);
        } else {
            runSequential(notifier);
        }
    }

    private void runSequential(RunNotifier notifier) {
        for (Runner runner : runners) {
            runner.run(notifier);
        }
    }

    private void runTestsInParallel(final RunNotifier notifier) {
        final int parrellelBits = Runtime.getRuntime().availableProcessors();
        final ExecutorService executorService = Executors.newFixedThreadPool(parrellelBits);
        for (final Runner runner : runners) {
            executorService.submit(new Runnable() {
                public void run() {
                    runner.run(notifier);
                }
            });
        }
        executorService.shutdown();
        try {
            executorService.awaitTermination(1000, TimeUnit.SECONDS);
        } catch (InterruptedException e) {
            executorService.shutdownNow();
            Thread.currentThread().interrupt();
        }
    }

}