Improve TestCoroutineScope

kotlinx-coroutines-test/common/src/TestBuilders.kt

@@ -43,25 +43,16 @@ import kotlin.coroutines.*
  */
 @ExperimentalCoroutinesApi
 public fun runBlockingTest(context: CoroutineContext = EmptyCoroutineContext, testBody: suspend TestCoroutineScope.() -> Unit) {
-    val (safeContext, dispatcher) = context.checkTestScopeArguments()
-    val startingJobs = safeContext.activeJobs()
-    val scope = TestCoroutineScope(safeContext)
+    val scope = TestCoroutineScope(context)
+    val scheduler = scope.coroutineContext[TestCoroutineScheduler]!!
     val deferred = scope.async {
         scope.testBody()
     }
-    dispatcher.scheduler.advanceUntilIdle()
+    scheduler.advanceUntilIdle()
     deferred.getCompletionExceptionOrNull()?.let {
         throw it
     }
     scope.cleanupTestCoroutines()
-    val endingJobs = safeContext.activeJobs()
-    if ((endingJobs - startingJobs).isNotEmpty()) {
-        throw UncompletedCoroutinesError("Test finished with active jobs: $endingJobs")
-    }
-}
-
-private fun CoroutineContext.activeJobs(): Set<Job> {
-    return checkNotNull(this[Job]).children.filter { it.isActive }.toSet()
 }
 
 /**
@@ -78,33 +69,3 @@ public fun TestCoroutineScope.runBlockingTest(block: suspend TestCoroutineScope.
 @ExperimentalCoroutinesApi
 public fun TestCoroutineDispatcher.runBlockingTest(block: suspend TestCoroutineScope.() -> Unit): Unit =
     runBlockingTest(this, block)
-
-internal fun CoroutineContext.checkTestScopeArguments(): Pair<CoroutineContext, TestDispatcher> {
-    val scheduler: TestCoroutineScheduler
-    val dispatcher = when (val dispatcher = get(ContinuationInterceptor)) {
-        is TestDispatcher -> {
-            val ctxScheduler = get(TestCoroutineScheduler)
-            if (ctxScheduler == null) {
-                scheduler = dispatcher.scheduler
-            } else {
-                require(dispatcher.scheduler === ctxScheduler) {
-                    "Both a TestCoroutineScheduler $ctxScheduler and TestDispatcher $dispatcher linked to " +
-                        "another scheduler were passed."
-                }
-                scheduler = ctxScheduler
-            }
-            dispatcher
-        }
-        null -> {
-            scheduler = get(TestCoroutineScheduler) ?: TestCoroutineScheduler()
-            TestCoroutineDispatcher(scheduler)
-        }
-        else -> throw IllegalArgumentException("Dispatcher must implement TestDispatcher: $dispatcher")
-    }
-    val exceptionHandler = get(CoroutineExceptionHandler).run {
-        this?.let { require(this is UncaughtExceptionCaptor) { "coroutineExceptionHandler must implement UncaughtExceptionCaptor: $this" } }
-        this ?: TestCoroutineExceptionHandler()
-    }
-    val job = get(Job) ?: SupervisorJob()
-    return Pair(this + scheduler + dispatcher + exceptionHandler + job, dispatcher)
-}

kotlinx-coroutines-test/common/src/TestCoroutineScope.kt

@@ -13,12 +13,12 @@ import kotlin.coroutines.*
 @ExperimentalCoroutinesApi
 public sealed interface TestCoroutineScope: CoroutineScope, UncaughtExceptionCaptor {
     /**
-     * Call after the test completes.
+     * Called after the test completes.
+     *
      * Calls [UncaughtExceptionCaptor.cleanupTestCoroutinesCaptor] and [DelayController.cleanupTestCoroutines].
      *
      * @throws Throwable the first uncaught exception, if there are any uncaught exceptions.
-     * @throws AssertionError if any pending tasks are active, however it will not throw for suspended
-     * coroutines.
+     * @throws AssertionError if any pending tasks are active.
      */
     @ExperimentalCoroutinesApi
     public fun cleanupTestCoroutines()
@@ -30,31 +30,81 @@ public sealed interface TestCoroutineScope: CoroutineScope, UncaughtExceptionCap
     public val testScheduler: TestCoroutineScheduler
 }
 
-private class TestCoroutineScopeImpl (
-    override val coroutineContext: CoroutineContext,
-    override val testScheduler: TestCoroutineScheduler
+private class TestCoroutineScopeImpl(
+    override val coroutineContext: CoroutineContext
 ):
     TestCoroutineScope,
     UncaughtExceptionCaptor by coroutineContext.uncaughtExceptionCaptor
 {
+    override val testScheduler: TestCoroutineScheduler
+        get() = coroutineContext[TestCoroutineScheduler]!!
+
+    /** These jobs existed before the coroutine scope was used, so it's alright if they don't get cancelled. */
+    private val initialJobs = coroutineContext.activeJobs()
+
     override fun cleanupTestCoroutines() {
         coroutineContext.uncaughtExceptionCaptor.cleanupTestCoroutinesCaptor()
         coroutineContext.delayController?.cleanupTestCoroutines()
+        val jobs = coroutineContext.activeJobs()
+        if ((jobs - initialJobs).isNotEmpty())
+            throw UncompletedCoroutinesError("Test finished with active jobs: $jobs")
     }
 }
 
+private fun CoroutineContext.activeJobs(): Set<Job> {
+    return checkNotNull(this[Job]).children.filter { it.isActive }.toSet()
+}
+
 /**
- * A scope which provides detailed control over the execution of coroutines for tests.
+ * A coroutine scope for launching test coroutines.
  *
- * If the provided context does not provide a [ContinuationInterceptor] (Dispatcher) or [CoroutineExceptionHandler], the
- * scope adds [TestCoroutineDispatcher] and [TestCoroutineExceptionHandler] automatically.
+ * It ensures that all the test module machinery is properly initialized.
+ * * If [context] doesn't define a [TestCoroutineScheduler] for orchestrating the virtual time used for delay-skipping,
+ *   a new one is created, unless a [TestDispatcher] is provided, in which case [TestDispatcher.scheduler] is used.
+ * * If [context] doesn't have a [ContinuationInterceptor], a [TestCoroutineDispatcher] is created.
+ * * If [context] does not provide a [CoroutineExceptionHandler], [TestCoroutineExceptionHandler] is created
+ *   automatically.
+ * * If [context] provides a [Job], that job is used for the new scope; otherwise, a [CompletableJob] is created.
  *
- * @param context an optional context that MAY provide [UncaughtExceptionCaptor] and/or [DelayController]
+ * @throws IllegalArgumentException if [context] has both [TestCoroutineScheduler] and a [TestDispatcher] linked to a
+ * different scheduler.
+ * @throws IllegalArgumentException if [context] has a [ContinuationInterceptor] that is not a [TestDispatcher].
+ * @throws IllegalArgumentException if [context] has an [CoroutineExceptionHandler] that is not an
+ * [UncaughtExceptionCaptor].
  */
 @Suppress("FunctionName")
 @ExperimentalCoroutinesApi
-public fun TestCoroutineScope(context: CoroutineContext = EmptyCoroutineContext): TestCoroutineScope =
-    context.checkTestScopeArguments().let { TestCoroutineScopeImpl(it.first, it.second.scheduler) }
+public fun TestCoroutineScope(context: CoroutineContext = EmptyCoroutineContext): TestCoroutineScope {
+    val scheduler: TestCoroutineScheduler
+    val dispatcher = when (val dispatcher = context[ContinuationInterceptor]) {
+        is TestDispatcher -> {
+            scheduler = dispatcher.scheduler
+            val ctxScheduler = context[TestCoroutineScheduler]
+            if (ctxScheduler != null) {
+                require(dispatcher.scheduler === ctxScheduler) {
+                    "Both a TestCoroutineScheduler $ctxScheduler and TestDispatcher $dispatcher linked to " +
+                        "another scheduler were passed."
+                }
+            }
+            dispatcher
+        }
+        null -> {
+            scheduler = context[TestCoroutineScheduler] ?: TestCoroutineScheduler()
+            TestCoroutineDispatcher(scheduler)
+        }
+        else -> throw IllegalArgumentException("Dispatcher must implement TestDispatcher: $dispatcher")
+    }
+    val exceptionHandler = context[CoroutineExceptionHandler].run {
+        this?.let {
+            require(this is UncaughtExceptionCaptor) {
+                "coroutineExceptionHandler must implement UncaughtExceptionCaptor: $context"
+            }
+        }
+        this ?: TestCoroutineExceptionHandler()
+    }
+    val job: Job = context[Job] ?: SupervisorJob()
+    return TestCoroutineScopeImpl(context + scheduler + dispatcher + exceptionHandler + job)
+}
 
 private inline val CoroutineContext.uncaughtExceptionCaptor: UncaughtExceptionCaptor
     get() {

kotlinx-coroutines-test/common/test/TestBuildersTest.kt

@@ -104,7 +104,7 @@ class TestBuildersTest {
     }
 
     @Test
-    fun whenInrunBlocking_runBlockingTest_nestsProperly() {
+    fun whenInRunBlocking_runBlockingTest_nestsProperly() {
         // this is not a supported use case, but it is possible so ensure it works
 
         val scope = TestCoroutineScope()

kotlinx-coroutines-test/common/test/TestCoroutineSchedulerTest.kt

@@ -5,6 +5,7 @@
 package kotlinx.coroutines.test
 
 import kotlinx.coroutines.*
+import kotlin.coroutines.*
 import kotlin.test.*
 
 class TestCoroutineSchedulerTest {

kotlinx-coroutines-test/common/test/TestCoroutineScopeTest.kt

@@ -50,6 +50,51 @@ class TestCoroutineScopeTest {
         }
     }
 
+    /** Tests that the cleanup procedure throws if there were uncompleted delays by the end. */
+    @Test
+    fun testPresentDelaysThrowing() {
+        val scope = TestCoroutineScope()
+        var result = false
+        scope.launch {
+            delay(5)
+            result = true
+        }
+        assertFalse(result)
+        assertFailsWith<AssertionError> { scope.cleanupTestCoroutines() }
+        assertFalse(result)
+    }
+
+    /** Tests that the cleanup procedure throws if there were active jobs by the end. */
+    @Test
+    fun testActiveJobsThrowing() {
+        val scope = TestCoroutineScope()
+        var result = false
+        val deferred = CompletableDeferred<String>()
+        scope.launch {
+            deferred.await()
+            result = true
+        }
+        assertFalse(result)
+        assertFailsWith<AssertionError> { scope.cleanupTestCoroutines() }
+        assertFalse(result)
+    }
+
+    /** Tests that the cleanup procedure doesn't throw if it detects that the job is already cancelled. */
+    @Test
+    fun testCancelledDelaysNotThrowing() {
+        val scope = TestCoroutineScope()
+        var result = false
+        val deferred = CompletableDeferred<String>()
+        val job = scope.launch {
+            deferred.await()
+            result = true
+        }
+        job.cancel()
+        assertFalse(result)
+        scope.cleanupTestCoroutines()
+        assertFalse(result)
+    }
+
     private val invalidContexts = listOf(
         Dispatchers.Default, // not a [TestDispatcher]
         TestCoroutineDispatcher() + TestCoroutineScheduler(), // the dispatcher is not linked to the scheduler