OkHttpIdlingResource + test rule for espresso tests

OkHttpIdlingResource.kt

package test.framework.idlingresources
import androidx.test.espresso.IdlingResource
import okhttp3.Dispatcher

/**
 * An [IdlingResource] for OkHttp, roughly based on [https://github.com/JakeWharton/okhttp-idling-resource], but in Kotlin,
 * and with a check to prevent experiencing this issue: ([https://github.com/JakeWharton/okhttp-idling-resource/issues/10]).
 */
class OkHttpIdlingResource constructor(private val name: String, private val dispatcher: Dispatcher) : IdlingResource {

    override fun getName(): String = name

    override fun isIdleNow(): Boolean = dispatcher.runningCallsCount() == 0

    @Volatile
    internal var callback: IdlingResource.ResourceCallback? = null

    init {
        check(dispatcher.idleCallback == null) {
            "dispatcher.idleCallback already set! Can't set for '$name' Perhaps the dispatcher is being shared (so you don't need this one)? ${dispatcher.idleCallback}"
        }
        dispatcher.idleCallback = Runnable { callback?.onTransitionToIdle() }
    }

    override fun registerIdleTransitionCallback(callback: IdlingResource.ResourceCallback) {
        this.callback = callback
    }

    fun cleanUp() {
        dispatcher.idleCallback = null
    }
}

OkHttpIdlingResourceRule.kt

package test.framework.rules

import androidx.test.espresso.IdlingRegistry
import okhttp3.Dispatcher
import org.junit.rules.TestWatcher
import org.junit.runner.Description
import test.framework.idlingresources.OkHttpIdlingResource


/**
 * A JUnit test rule to register an [androidx.test.espresso.IdlingResource] for an OkHttp [Dispatcher].
 * Typical usage: OkHttpIdlingResourceRule("OkHttp dispatcher") { okHttpClient.dispatcher }.
 * It lazily fetches the dispatcher at the start of the test rather than at rule declaration, so it will be ready.
 */
class OkHttpIdlingResourceRule(private val name: String, private val dispatcher: () -> Dispatcher) : TestWatcher() {
    private var idlingResource: OkHttpIdlingResource? = null

    override fun starting(description: Description?) {
        idlingResource = OkHttpIdlingResource("OkHttp ($name)", dispatcher())
        IdlingRegistry.getInstance().register(idlingResource)
    }

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