Cancel the child coroutines if the test body has thrown

Fixes #1910

Author: dkhalanskyjb

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

@@ -141,6 +141,18 @@ public expect class TestResult
  *
  * ### Failures
  *
+ * #### Test body failures
+ *
+ * If the test body finishes with an exception, then this exception will be thrown at the end of the test. Additionally,
+ * to prevent child coroutines getting stuck, the whole scope will be cancelled in this case.
+ *
+ * #### Reported exceptions
+ *
+ * Exceptions reported to the test coroutine scope via [TestCoroutineScope.reportException] will be thrown at the end.
+ * By default (without passing an explicit [TestExceptionHandler]), this includes all unhandled exceptions.
+ *
+ * #### Uncompleted coroutines
+ *
  * This method requires that all coroutines launched inside [testBody] complete, or are cancelled. Otherwise, the test
  * will be failed (which, on JVM and Native, means that [runTest] itself will throw [AssertionError],
  * whereas on JS, the `Promise` will fail with it).
@@ -151,8 +163,6 @@ public expect class TestResult
  * idle before throwing [AssertionError]. If some dispatcher linked to [TestCoroutineScheduler] receives a
  * task during that time, the timer gets reset.
  *
- * Unhandled exceptions thrown by coroutines in the test will be rethrown at the end of the test.
- *
  * ### Configuration
  *
  * [context] can be used to affect the environment of the code under test. Beside just being passed to the coroutine
@@ -178,7 +188,23 @@ public fun runTest(
         }
         var completed = false
         while (!completed) {
-            scheduler.advanceUntilIdle()
+            while (scheduler.tryRunNextTask()) {
+                if (deferred.isCompleted && deferred.getCompletionExceptionOrNull() != null && testScope.isActive) {
+                    /**
+                     * Here, we already know how the test will finish: it will throw
+                     * [Deferred.getCompletionExceptionOrNull]. Therefore, we won't care if there are uncompleted jobs,
+                     * and may as well just exit right here. However, in order to lower the surprise factor, we
+                     * cancel the child jobs here and wait for them to finish instead of dropping them: there could be
+                     * some cleanup procedures involved, and not having finalizers run could mean leaking resources.
+                     *
+                     * Another approach to take if this turns out not to be enough and some child jobs still fail is to
+                     * only make at most a fixed number of [TestCoroutineScheduler.tryRunNextTask] once we detect the
+                     * failure with which the test will finish. This has the downside that there is still some
+                     * negligible risk of not running the finalizers.
+                     */
+                    testScope.cancel()
+                }
+            }
             if (deferred.isCompleted) {
                 /* don't even enter `withTimeout`; this allows to use a timeout of zero to check that there are no
                    non-trivial dispatches. */

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

@@ -81,6 +81,21 @@ public class TestCoroutineScheduler : AbstractCoroutineContextElement(TestCorout
         }
     }
 
+    /**
+     * Runs the next enqueued task, advancing the virtual time to the time of its scheduled awakening.
+     */
+    internal fun tryRunNextTask(): Boolean {
+        val event = synchronized(lock) {
+            val event = events.removeFirstOrNull() ?: return false
+            if (currentTime > event.time)
+                currentTimeAheadOfEvents()
+            currentTime = event.time
+            event
+        }
+        event.dispatcher.processEvent(event.time, event.marker)
+        return true
+    }
+
     /**
      * Runs the enqueued tasks in the specified order, advancing the virtual time as needed until there are no more
      * tasks associated with the dispatchers linked to this scheduler.
@@ -92,14 +107,7 @@ public class TestCoroutineScheduler : AbstractCoroutineContextElement(TestCorout
     @ExperimentalCoroutinesApi
     public fun advanceUntilIdle() {
         while (!events.isEmpty) {
-            val event = synchronized(lock) {
-                val event = events.removeFirstOrNull() ?: return
-                if (currentTime > event.time)
-                    currentTimeAheadOfEvents()
-                currentTime = event.time
-                event
-            }
-            event.dispatcher.processEvent(event.time, event.marker)
+            tryRunNextTask()
         }
     }
 

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

@@ -153,4 +153,54 @@ class RunTestTest {
         }
     }
 
+    @Test
+    fun reproducer2405() = runTest {
+        val dispatcher = StandardTestDispatcher(testScheduler)
+        var collectedError = false
+        withContext(dispatcher) {
+            flow { emit(1) }
+                .combine(
+                    flow<String> { throw IllegalArgumentException() }
+                ) { int, string -> int.toString() + string }
+                .catch { emit("error") }
+                .collect {
+                    assertEquals("error", it)
+                    collectedError = true
+                }
+        }
+        assertTrue(collectedError)
+    }
+
+    /** Tests that, once the test body has thrown, the child coroutines are cancelled. */
+    @Test
+    fun testChildrenCancellationOnTestBodyFailure() {
+        var job: Job? = null
+        testResultMap({
+            assertFailsWith<AssertionError> { it() }
+            assertTrue(job!!.isCancelled)
+        }, {
+            runTest {
+                job = launch {
+                    while (true) {
+                        delay(1000)
+                    }
+                }
+                throw AssertionError()
+            }
+        })
+    }
+
+    /** Tests that [runTest] reports [TimeoutCancellationException]. */
+    @Test
+    fun testTimeout() = testResultMap({
+        assertFailsWith<TimeoutCancellationException> { it() }
+    }, {
+        runTest {
+            withTimeout(50) {
+                launch {
+                    delay(1000)
+                }
+            }
+        }
+    })
 }