package uk.co.ribot.androidboilerplate.util;

import org.junit.rules.TestRule;
import org.junit.runner.Description;
import org.junit.runners.model.Statement;

import rx.Scheduler;
import rx.android.plugins.RxAndroidPlugins;
import rx.android.plugins.RxAndroidSchedulersHook;
import rx.functions.Func1;
import rx.plugins.RxJavaHooks;
import rx.schedulers.Schedulers;

/**
 * This rule registers SchedulerHooks for RxJava and RxAndroid to ensure that subscriptions
 * always subscribeOn and observeOn Schedulers.immediate().
 * Warning, this rule will reset RxAndroidPlugins and RxJavaPlugins before and after each test so
 * if the application code uses RxJava plugins this may affect the behaviour of the testing method.
 */
public class RxSchedulersOverrideRule implements TestRule {

    private final RxAndroidSchedulersHook mRxAndroidSchedulersHook = new RxAndroidSchedulersHook() {
        @Override
        public Scheduler getMainThreadScheduler() {
            return Schedulers.immediate();
        }
    };

    private final Func1<Scheduler, Scheduler> mRxJavaImmediateScheduler =
            new Func1<Scheduler, Scheduler>() {
                @Override
                public Scheduler call(Scheduler scheduler) {
                    return Schedulers.immediate();
                }
            };

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

                RxJavaHooks.reset();
                RxJavaHooks.setOnIOScheduler(mRxJavaImmediateScheduler);
                RxJavaHooks.setOnNewThreadScheduler(mRxJavaImmediateScheduler);

                base.evaluate();

                RxAndroidPlugins.getInstance().reset();
                RxJavaHooks.reset();
            }
        };
    }
}

@@ -151,25 +151,73 @@
 tement;%0A
-%0Aimport r
+import java.util.concurrent.Callable;%0A%0Aimport io.reactive
 x.Schedu
@@ -220,33 +220,43 @@
 heduler;%0Aimport 
-r
+io.reactive
 x.android.plugin
@@ -286,60 +286,60 @@
 ort 
-rx.android.plugins.RxAndroidSchedulersHook;%0Aimport r
+io.reactivex.annotations.NonNull;%0Aimport io.reactive
 x.fu
@@ -350,17 +350,20 @@
 ons.Func
-1
+tion
 ;%0Aimport
@@ -363,17 +363,27 @@
 %0Aimport 
-r
+io.reactive
 x.plugin
@@ -390,20 +390,22 @@
 s.RxJava
-Hook
+Plugin
 s;%0Aimpor
@@ -406,17 +406,27 @@
 %0Aimport 
-r
+io.reactive
 x.schedu
@@ -474,21 +474,15 @@
 ers 
-SchedulerHook
+Handler
 s fo
@@ -534,17 +534,17 @@
 ptions%0A 
-*
++
  always 
@@ -580,24 +580,25 @@
 edulers.
-immediat
+trampolin
 e().%0A * 
@@ -875,25 +875,46 @@
 e final 
-RxAndroid
+Function%3CCallable%3CScheduler%3E, 
 Schedule
@@ -914,21 +914,17 @@
 cheduler
-sHook
+%3E
  mRxAndr
@@ -946,49 +946,84 @@
 ok =
- new RxAndroidSchedulersHook() %7B%0A
+%0A            new Function%3CCallable%3CScheduler%3E, Scheduler%3E() %7B%0A      
         
 @Ove
@@ -1010,32 +1010,34 @@
 %7B%0A              
+  
 @Override%0A      
@@ -1022,32 +1022,40 @@
       @Override%0A
+        
         public S
@@ -1067,35 +1067,113 @@
 ler 
-getMainThreadScheduler() %7B%0A
+apply(@NonNull Callable%3CScheduler%3E schedulerCallable)%0A                        throws Exception %7B%0A        
     
@@ -1198,37 +1198,45 @@
 edulers.
-immediate();%0A
+trampoline();%0A       
         
 %7D%0A    %7D;
@@ -1223,26 +1223,35 @@
             
+ 
 %7D%0A
+        
     %7D;%0A%0A    
@@ -1264,25 +1264,28 @@
 e final Func
-1
+tion
 %3CScheduler, 
@@ -1347,9 +1347,12 @@
 Func
-1
+tion
 %3CSch
@@ -1437,13 +1437,23 @@
 ler 
-call(
+apply(@NonNull 
 Sche
@@ -1468,16 +1468,33 @@
 heduler)
+ throws Exception
  %7B%0A     
@@ -1530,16 +1530,17 @@
 ers.
-immediat
+trampolin
 e();
@@ -1802,38 +1802,24 @@
 roidPlugins.
-getInstance().
 reset();%0A   
@@ -1852,30 +1852,25 @@
 ins.
-g
+s
 etIn
-stance().register
+itMainThread
 Sche
@@ -1874,21 +1874,23 @@
 cheduler
-sHook
+Handler
 (mRxAndr
@@ -1924,36 +1924,38 @@
           RxJava
-Hook
+Plugin
 s.reset();%0A     
@@ -1975,21 +1975,21 @@
 Java
-Hook
+Plugin
 s.set
-OnIO
+Io
 Sche
@@ -1985,32 +1985,39 @@
 s.setIoScheduler
+Handler
 (mRxJavaImmediat
@@ -2055,19 +2055,19 @@
 Java
-Hook
+Plugin
 s.set
-On
 NewT
@@ -2080,16 +2080,23 @@
 cheduler
+Handler
 (mRxJava
@@ -2188,22 +2188,8 @@
 ins.
-getInstance().
 rese
@@ -2215,20 +2215,22 @@
   RxJava
-Hook
+Plugin
 s.reset(
@@ -2263,10 +2263,11 @@
 ;%0A    %7D%0A
+%0A
 %7D%0A

Migrate RxSchedulersOverrideRule to RxJava2.