Fix 'TestCoroutineScope.runTest' not working

Author: dkhalanskyjb

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

@@ -239,10 +239,8 @@ internal expect fun createTestResult(testProcedure: suspend () -> Unit): TestRes
 /**
  * Runs a test in a [TestCoroutineScope] based on this one.
  *
- * Calls [runTest] using a coroutine context from this [TestCoroutineScope]. The [TestCoroutineScope] used to run
- * [block] will be different from this one, but will use its [Job] as a parent; therefore, even if calling
- * [TestCoroutineScope.cleanupTestCoroutines] on this scope were to complete its job, [runTest] won't complete it at the
- * end of the test.
+ * Calls [runTest] using a coroutine context from this [TestCoroutineScope]. The [TestCoroutineScope] used to run the
+ * [block] will be different from this one, but will use its [Job] as a parent.
  *
  * Since this function returns [TestResult], in order to work correctly on the JS, its result must be returned
  * immediately from the test body. See the docs for [TestResult] for details.

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

@@ -179,18 +179,22 @@ public fun createTestCoroutineScope(context: CoroutineContext = EmptyCoroutineCo
         else -> throw IllegalArgumentException("Dispatcher must implement TestDispatcher: $dispatcher")
     }
     var scope: TestCoroutineScopeImpl? = null
-    val exceptionHandler = when (val exceptionHandler = context[CoroutineExceptionHandler]) {
-        is UncaughtExceptionCaptor -> exceptionHandler
-        null -> {
-            val lock = SynchronizedObject()
-            CoroutineExceptionHandler { _, throwable ->
+    val ownExceptionHandler = run {
+        val lock = SynchronizedObject()
+        object : AbstractCoroutineContextElement(CoroutineExceptionHandler), TestCoroutineScopeExceptionHandler {
+            override fun handleException(context: CoroutineContext, exception: Throwable) {
                 val reported = synchronized(lock) {
-                    scope!!.reportException(throwable)
+                    scope!!.reportException(exception)
                 }
                 if (!reported)
-                    throw throwable // let this exception crash everything
+                    throw exception // let this exception crash everything
             }
         }
+    }
+    val exceptionHandler = when (val exceptionHandler = context[CoroutineExceptionHandler]) {
+        is UncaughtExceptionCaptor -> exceptionHandler
+        null -> ownExceptionHandler
+        is TestCoroutineScopeExceptionHandler -> ownExceptionHandler
         else -> throw IllegalArgumentException(
             "A CoroutineExceptionHandler was passed to TestCoroutineScope. " +
                 "Please pass it as an argument to a `launch` or `async` block on an already-created scope " +
@@ -203,6 +207,10 @@ public fun createTestCoroutineScope(context: CoroutineContext = EmptyCoroutineCo
     }
 }
 
+/** A marker that shows that this [CoroutineExceptionHandler] was created for [TestCoroutineScope]. With this,
+ * constructing a new [TestCoroutineScope] with the [CoroutineScope.coroutineContext] of an existing one will override
+ * the exception handler, instead of failing. */
+private interface TestCoroutineScopeExceptionHandler: CoroutineExceptionHandler
 
 private inline val CoroutineContext.delayController: DelayController?
     get() {

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

@@ -275,4 +275,22 @@ class RunTestTest {
             throw TestException("w")
         }
     })
+
+    /** Tests that [TestCoroutineScope.runTest] does not inherit the exception handler and works. */
+    @Test
+    fun testScopeRunTestExceptionHandler(): TestResult {
+        val scope = createTestCoroutineScope()
+        return testResultMap({
+            try {
+                it()
+                fail("should not be reached")
+            } catch (e: TestException) {
+                scope.cleanupTestCoroutines() // should not fail
+            }
+        }, {
+            scope.runTest {
+                launch(SupervisorJob()) { throw TestException("x") }
+            }
+        })
+    }
 }

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

@@ -153,6 +153,32 @@ class TestCoroutineScopeTest {
         }
     }
 
+    /** Tests that constructing a new [TestCoroutineScope] using another one's scope works and overrides the exception
+     * handler. */
+    @Test
+    fun testCopyingContexts() {
+        val deferred = CompletableDeferred<Unit>()
+        val scope1 = createTestCoroutineScope()
+        scope1.launch { deferred.await() } // a pending job in the outer scope
+        val scope2 = createTestCoroutineScope(scope1.coroutineContext)
+        val scope3 = createTestCoroutineScope(scope1.coroutineContext)
+        assertEquals(
+            scope1.coroutineContext.minusKey(CoroutineExceptionHandler),
+            scope2.coroutineContext.minusKey(CoroutineExceptionHandler))
+        scope2.launch(SupervisorJob()) { throw TestException("x") } // will fail the cleanup of scope2
+        try {
+            scope2.cleanupTestCoroutines()
+            fail("should not be reached")
+        } catch (e: TestException) { }
+        scope3.cleanupTestCoroutines() // the pending job in the outer scope will not cause this to fail
+        try {
+            scope1.cleanupTestCoroutines()
+            fail("should not be reached")
+        } catch (e: UncompletedCoroutinesError) {
+            // the pending job in the outer scope
+        }
+    }
+
     companion object {
         internal val invalidContexts = listOf(
             Dispatchers.Default, // not a [TestDispatcher]