A JUnit 4 Rule and JUnit 5 Extension for utilizing TestCoroutineDispatcher and TestCoroutineScope from kotlinx.coroutines-test

TestCoroutineExtension.kt

import kotlinx.coroutines.Dispatchers
import kotlinx.coroutines.test.TestCoroutineDispatcher
import kotlinx.coroutines.test.TestCoroutineScope
import kotlinx.coroutines.test.resetMain
import kotlinx.coroutines.test.setMain
import org.junit.jupiter.api.extension.AfterAllCallback
import org.junit.jupiter.api.extension.AfterEachCallback
import org.junit.jupiter.api.extension.BeforeAllCallback
import org.junit.jupiter.api.extension.ExtendWith
import org.junit.jupiter.api.extension.ExtensionContext
import org.junit.jupiter.api.extension.TestInstancePostProcessor

@ExtendWith(TestCoroutineExtension::class)
interface CoroutineTest {
  var testScope: TestCoroutineScope
  var dispatcher: TestCoroutineDispatcher
}

/**
 * JUnit 5 Extension for automatically creating a [TestCoroutineDispatcher],
 * then a [TestCoroutineScope] with the same CoroutineContext.
 *
 * [TestCoroutineScope.cleanupTestCoroutines] is called in afterEach
 * instead of afterAll in case Lifecycle.PER_CLASS is selected,
 * and will cause an exception if any coroutines are leaked.
 *
 * Usage of an extension in a Kotlin JUnit 5 test:
 *
 * class MyTest : CoroutineTest {
 *
 *   override lateinit var testScope: TestCoroutineScope
 *   override lateinit var dispatcher: TestCoroutineDispatcher
 * }
 *
 */
@ExperimentalCoroutinesApi
class TestCoroutineExtension : TestInstancePostProcessor, BeforeAllCallback, AfterEachCallback, AfterAllCallback {

  val dispatcher = TestCoroutineDispatcher()
  val testScope = TestCoroutineScope(dispatcher)

  override fun postProcessTestInstance(testInstance: Any?, context: ExtensionContext?) {

    (testInstance as? CoroutineTest)?.let { coroutineTest ->

      coroutineTest.testScope = testScope
      coroutineTest.dispatcher = dispatcher
    }
  }

  override fun beforeAll(context: ExtensionContext?) {
    Dispatchers.setMain(dispatcher)
  }

  override fun afterEach(context: ExtensionContext?) {
    testScope.cleanupTestCoroutines()
  }

  override fun afterAll(context: ExtensionContext?) {
    Dispatchers.resetMain()
  }

}

TestCoroutineRule.kt

import kotlinx.coroutines.CoroutineDispatcher
import kotlinx.coroutines.Dispatchers
import kotlinx.coroutines.ExperimentalCoroutinesApi
import kotlinx.coroutines.test.TestCoroutineDispatcher
import kotlinx.coroutines.test.TestCoroutineScope
import kotlinx.coroutines.test.resetMain
import kotlinx.coroutines.test.setMain
import org.junit.rules.TestRule
import org.junit.runner.Description
import org.junit.runners.model.Statement
import kotlin.coroutines.ContinuationInterceptor

/**
 * JUnit 4 Rule for automatically creating a [TestCoroutineDispatcher],
 * then a [TestCoroutineScope] with the same CoroutineContext.
 *
 * [TestCoroutineScope.cleanupTestCoroutines] is called after execution of each test,
 * and will cause an exception if any coroutines are leaked.
 *
 * Usage of a rule in a Kotlin JUnit 4 test is:
 *
 * class MyTest {
 *
 *   @get:Rule val testRule = TestCoroutineRule()
 *
 *   val dispatcher = testRule.dispatcher
 * }
 *
 */
@ExperimentalCoroutinesApi
class TestCoroutineRule : TestRule, TestCoroutineScope by TestCoroutineScope() {

  val dispatcher = coroutineContext[ContinuationInterceptor] as TestCoroutineDispatcher

  override fun apply(base: Statement, description: Description): Statement {
    return object : Statement() {
      @Throws(Throwable::class)
      override fun evaluate() {

        Dispatchers.setMain(dispatcher)

        // everything above this happens before the test
        base.evaluate()
        // everything below this happens after the test

        cleanupTestCoroutines()
        Dispatchers.resetMain()
      }
    }
  }
}