MainCoroutineRule.kt

class MainCoroutineRule : TestWatcher() {

    private val testDispatcher = TestCoroutineDispatcher()

    override fun starting(description: Description?) {
        Dispatchers.setMain(testDispatcher)
    }

    fun setup(setupBody: () -> Unit){
        testDispatcher.pauseDispatcher()
        setupBody.invoke()
    }

    fun run(testBody: () -> Unit){
        testDispatcher.resumeDispatcher()
        testBody.invoke()
    }

    override fun finished(description: Description?) {
        Dispatchers.resetMain()
        testDispatcher.cleanupTestCoroutines()
    }
}