RxJava2 Test rule

ImmediateRxSchedulersOverrideRule.java

/**
 * This rule registers SchedulerHooks for RxJava and RxAndroid to ensure that subscriptions
 * always subscribeOn and observeOn Schedulers.immediate().
 * Warning, this rule will resetProcedureStatus RxAndroidPlugins and RxJavaPlugins after each test so
 * if the application code uses RxJava plugins this may affect the behaviour of the testing method.
 * <p>
 *
 * See https://medium.com/@fabioCollini/testing-asynchronous-rxjava-code-using-mockito-8ad831a16877#.ahj5h7jmg
 * See https://github.com/fabioCollini/TestingRxJavaUsingMockito/blob/master/app/src/test/java/it/codingjam/testingrxjava/TestSchedulerRule.java
 */
public class ImmediateRxSchedulersOverrideRule implements TestRule {

  private Scheduler SCHEDULER_INSTANCE = Schedulers.trampoline();

  private Function<Scheduler, Scheduler> schedulerFunction = new Function<Scheduler, Scheduler>() {
    @Override
    public Scheduler apply(Scheduler scheduler) throws Exception {
      return SCHEDULER_INSTANCE;
    }
  };

  private Function<Callable<Scheduler>, Scheduler> schedulerFunctionLazy = new Function<Callable<Scheduler>, Scheduler>() {
    @Override
    public Scheduler apply(Callable<Scheduler> schedulerCallable) throws Exception {
      return SCHEDULER_INSTANCE;
    }
  };

  @Override
  public Statement apply(final Statement base, Description description) {
    return new Statement() {
      @Override
      public void evaluate() throws Throwable {
        RxAndroidPlugins.reset();
        RxAndroidPlugins.setInitMainThreadSchedulerHandler(schedulerFunctionLazy);

        RxJavaPlugins.reset();
        RxJavaPlugins.setIoSchedulerHandler(schedulerFunction);
        RxJavaPlugins.setNewThreadSchedulerHandler(schedulerFunction);
        RxJavaPlugins.setComputationSchedulerHandler(schedulerFunction);

        base.evaluate();

        RxAndroidPlugins.reset();
        RxJavaPlugins.reset();
      }
    };
  }
}