Coroutine dispatcher idling resources for espresso tests

DispatchersIdlingResourceRule.kt

package test.framework.rules

import androidx.test.espresso.IdlingRegistry
import org.junit.rules.TestWatcher
import org.junit.runner.Description
import test.framework.idlingresources.DispatcherWithIdlingResource

/**
 * A JUnit test rule that registers an idling resource for each [DispatcherWithIdlingResource] given
 *
 */
class DispatchersIdlingResourceRule(private vararg val dispatchers: DispatcherWithIdlingResource) : TestWatcher() {
    override fun starting(description: Description?) {
        dispatchers.forEach {
            IdlingRegistry.getInstance().register(it.idlingResource)
        }
    }

    override fun finished(description: Description?) {
        dispatchers.forEach {
            IdlingRegistry.getInstance().unregister(it.idlingResource)
        }
    }
}

DispatcherWithIdlingResource.kt

package test.framework.idlingresources

import androidx.test.espresso.IdlingResource

/**
 * A shared interface for our wrappers around [kotlinx.coroutines.CoroutineDispatcher] and 
 * [kotlinx.coroutines.MainCoroutineDispatcher] that expose an [IdlingResource] for each of them.
 */
interface DispatcherWithIdlingResource {
    val idlingResource: IdlingResource
}

EspressoDispatchers.kt

package com.themightyjon.app.shared.coroutines

import androidx.test.espresso.idling.CountingIdlingResource
import kotlinx.coroutines.CoroutineDispatcher
import kotlinx.coroutines.Dispatchers
import kotlinx.coroutines.MainCoroutineDispatcher
import test.framework.idlingresources.HasIdlingResource
import kotlin.coroutines.CoroutineContext

/*
NOTE this file should go into the androidTest sources, not main. In tests you can substitute the dispatchers that 
you inject -- you do inject your dispatchers don't you? ;) -- with these ones
*/

class EspressoDispatchers : AppDispatchers {
    override val IO = EspressoTrackedDispatcher(Dispatchers.IO)
    override val Main = EspressoTrackedMainDispatcher(Dispatchers.Main)
}

fun delegateDispatchWithCounting(
    delegateDispatcher: CoroutineDispatcher,
    context: CoroutineContext,
    block: Runnable,
    idlingResource: CountingIdlingResource
) {
    idlingResource.increment()
    delegateDispatcher.dispatch(context, Runnable {
        try {
            block.run()
        } finally {
            idlingResource.decrement()
        }
    })
}

/**
 * Decorates [CoroutineDispatcher] adding a [CountingIdlingResource]. Based on [https://github.com/Kotlin/kotlinx.coroutines/issues/242].
 */
class EspressoTrackedDispatcher(private val delegateDispatcher: CoroutineDispatcher) : CoroutineDispatcher(), DispatcherWithIdlingResource {
    override val idlingResource: CountingIdlingResource = CountingIdlingResource("EspressoTrackedDispatcher for $delegateDispatcher")

    override fun dispatch(context: CoroutineContext, block: Runnable) = delegateDispatchWithCounting(delegateDispatcher, context, block, idlingResource)
}

/**
 * Decorates [MainCoroutineDispatcher] adding a [CountingIdlingResource]. Based on [https://github.com/Kotlin/kotlinx.coroutines/issues/242].
 * The main dispatcher is a totally different class so we have to duplicate EspressoTrackedDispatcher to provide a dispatcher of that type too.
 */
class EspressoTrackedMainDispatcher(private val delegateDispatcher: MainCoroutineDispatcher) : MainCoroutineDispatcher(), DispatcherWithIdlingResource {
    override val idlingResource: CountingIdlingResource = CountingIdlingResource("EspressoTrackedMainDispatcher for $delegateDispatcher")

    override val immediate: MainCoroutineDispatcher =
        if (delegateDispatcher.immediate === delegateDispatcher) this else EspressoTrackedMainDispatcher(delegateDispatcher.immediate)

    override fun dispatch(context: CoroutineContext, block: Runnable) = delegateDispatchWithCounting(delegateDispatcher, context, block, idlingResource)
}

MyDispatchers.kt

package com.themightyjon.app.shared.coroutines

import kotlinx.coroutines.CoroutineDispatcher
import kotlinx.coroutines.MainCoroutineDispatcher

/*
NOTE this file should go into the main sources. Inject AppDispatchers wherever you need to reference dispatchers instead of 
referencing them statically. Then they will be substitutable in your tests.
*/

@Suppress("PropertyName") // Made to match Dispatchers in platform for ease of change
interface AppDispatchers {
    val IO: CoroutineDispatcher
    val Main: MainCoroutineDispatcher
}

class ProductionDispatchers : MyDispatchers {
    override val IO: CoroutineDispatcher = Dispatchers.IO
    override val Main: MainCoroutineDispatcher = Dispatchers.Main
}

Hi man @jonreeve
This is awesome, i was trying to implement this on my project but still doesn't work, currently still using a wait mechanism from androidx.test.espresso.UiController.loopMainThreadForAtLeast()

was try to use EspressoDispatchers combined with hilt but not working,
this is how i trying to implement:

@Inject
lateinit var dispatcher: DispatcherProvider //actually EspressoDispatchers injected as @Singleton instance

@Before
fun setUp() {
   hiltRule.inject(this)
   IdlingRegistry.getInstance().register((dispatcher.main as DispatcherWithIdlingResource).idlingResource)
   IdlingRegistry.getInstance().register((dispatcher.io as DispatcherWithIdlingResource).idlingResource)
}

@After
fun tearDown() {
   IdlingRegistry.getInstance().register((dispatcher.main as DispatcherWithIdlingResource).idlingResource)
   IdlingRegistry.getInstance().register((dispatcher.io as DispatcherWithIdlingResource).idlingResource)
}

Hi man @jonreeve This is awesome, i was trying to implement this on my project but still doesn't work, currently still using a wait mechanism from androidx.test.espresso.UiController.loopMainThreadForAtLeast()

was try to use EspressoDispatchers combined with hilt but not working, this is how i trying to implement:

@Inject
lateinit var dispatcher: DispatcherProvider //actually EspressoDispatchers injected as @Singleton instance

@Before
fun setUp() {
   hiltRule.inject(this)
   IdlingRegistry.getInstance().register((dispatcher.main as DispatcherWithIdlingResource).idlingResource)
   IdlingRegistry.getInstance().register((dispatcher.io as DispatcherWithIdlingResource).idlingResource)
}

@After
fun tearDown() {
   IdlingRegistry.getInstance().register((dispatcher.main as DispatcherWithIdlingResource).idlingResource)
   IdlingRegistry.getInstance().register((dispatcher.io as DispatcherWithIdlingResource).idlingResource)
}

I think that you need to actually replace the original dispatchers by the espresso dispatchers by hilt...

Example (I am using DispatchersModule to provide dispatchers and replacing the whole module for tests):

@Module
@TestInstallIn(
    components = [SingletonComponent::class],
    replaces = [DispatchersModule::class]
)
object TestDispatchersModule {

    @Provides
    @Singleton
    fun provideTestCoroutineDispatchers(testDispatchersProvider: TestDispatchersProvider): CoroutineDispatchers =
        testDispatchersProvider.provideTestCoroutineDispatchers()
}