Ensure runTest unsubscribes from the exception handler (#3909)

Fixes #3897

Author: dkhalanskyjb

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

@@ -239,6 +239,7 @@ internal class TestScopeImpl(context: CoroutineContext) :
             uncaughtExceptions
         }
         if (exceptions.isNotEmpty()) {
+            ExceptionCollector.removeOnExceptionCallback(lock)
             throw UncaughtExceptionsBeforeTest().apply {
                 for (e in exceptions)
                     addSuppressed(e)

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

@@ -47,14 +47,16 @@ fun testResultMap(block: (() -> Unit) -> Unit, test: () -> TestResult): TestResu
  */
 expect fun testResultChain(block: () -> TestResult, after: (Result<Unit>) -> TestResult): TestResult
 
-fun testResultChain(vararg chained: (Result<Unit>) -> TestResult): TestResult =
+fun testResultChain(vararg chained: (Result<Unit>) -> TestResult, initialResult: Result<Unit> = Result.success(Unit)): TestResult =
     if (chained.isEmpty()) {
-        createTestResult { }
+        createTestResult {
+            initialResult.getOrThrow()
+        }
     } else {
         testResultChain(block = {
-            chained[0](Result.success(Unit))
+            chained[0](initialResult)
         }) {
-            testResultChain(*chained.drop(1).toTypedArray())
+            testResultChain(*chained.drop(1).toTypedArray(), initialResult = it)
         }
     }
 

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

@@ -408,4 +408,69 @@ class RunTestTest {
     fun testCoroutineCompletingWithoutDispatch() = runTest(timeout = Duration.INFINITE) {
         launch(Dispatchers.Default) { delay(100) }
     }
+
+    /**
+     * Tests that [runTest] cleans up the exception handler even if it threw on initialization.
+     *
+     * This test must be run manually, because it writes garbage to the log.
+     *
+     * The JVM-only source set contains a test equivalent to this one that isn't ignored.
+     */
+    @Test
+    @Ignore
+    fun testExceptionCaptorCleanedUpOnPreliminaryExit(): TestResult = testResultChain({
+        // step 1: installing the exception handler
+        println("step 1")
+        runTest { }
+    }, {
+        it.getOrThrow()
+        // step 2: throwing an uncaught exception to be caught by the exception-handling system
+        println("step 2")
+        createTestResult {
+            launch(NonCancellable) { throw TestException("A") }
+        }
+    }, {
+        it.getOrThrow()
+        // step 3: trying to run a test should immediately fail, even before entering the test body
+        println("step 3")
+        try {
+            runTest {
+                fail("unreached")
+            }
+            fail("unreached")
+        } catch (e: UncaughtExceptionsBeforeTest) {
+            val cause = e.suppressedExceptions.single()
+            assertIs<TestException>(cause)
+            assertEquals("A", cause.message)
+        }
+        // step 4: trying to run a test again should not fail with an exception
+        println("step 4")
+        runTest {
+        }
+    }, {
+        it.getOrThrow()
+        // step 5: throwing an uncaught exception to be caught by the exception-handling system, again
+        println("step 5")
+        createTestResult {
+            launch(NonCancellable) { throw TestException("B") }
+        }
+    }, {
+        it.getOrThrow()
+        // step 6: trying to run a test should immediately fail, again
+        println("step 6")
+        try {
+            runTest {
+                fail("unreached")
+            }
+            fail("unreached")
+        } catch (e: Exception) {
+            val cause = e.suppressedExceptions.single()
+            assertIs<TestException>(cause)
+            assertEquals("B", cause.message)
+        }
+        // step 7: trying to run a test again should not fail with an exception, again
+        println("step 7")
+        runTest {
+        }
+    })
 }

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

@@ -495,9 +495,11 @@ class TestScopeTest {
      * Tests that the [TestScope] exception reporting mechanism will report the exceptions that happen between
      * different tests.
      *
-     * This test must be ran manually, because such exceptions still go through the global exception handler
+     * This test must be run manually, because such exceptions still go through the global exception handler
      * (as there's no guarantee that another test will happen), and the global exception handler will
      * log the exceptions or, on Native, crash the test suite.
+     *
+     * The JVM-only source set contains a test equivalent to this one that isn't ignored.
      */
     @Test
     @Ignore

kotlinx-coroutines-test/jvm/test/UncaughtExceptionsTest.kt

@@ -0,0 +1,50 @@
+/*
+ * Copyright 2016-2023 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license.
+ */
+package kotlinx.coroutines.test
+
+import org.junit.Test
+import kotlin.test.*
+
+/**
+ * Tests that check the behavior of the test framework when there are stray uncaught exceptions.
+ * These tests are JVM-only because only the JVM allows to set a global uncaught exception handler and validate the
+ * behavior without checking the test logs.
+ * Nevertheless, each test here has a corresponding test in the common source set that can be run manually.
+ */
+class UncaughtExceptionsTest {
+
+    val oldExceptionHandler = Thread.getDefaultUncaughtExceptionHandler()
+    val uncaughtExceptions = mutableListOf<Throwable>()
+
+    @BeforeTest
+    fun setUp() {
+        Thread.setDefaultUncaughtExceptionHandler { thread, throwable ->
+            uncaughtExceptions.add(throwable)
+        }
+    }
+
+    @AfterTest
+    fun tearDown() {
+        Thread.setDefaultUncaughtExceptionHandler(oldExceptionHandler)
+    }
+
+    @Test
+    fun testReportingStrayUncaughtExceptionsBetweenTests() {
+        TestScopeTest().testReportingStrayUncaughtExceptionsBetweenTests()
+        assertEquals(1, uncaughtExceptions.size, "Expected 1 uncaught exception, but got $uncaughtExceptions")
+        val exception = assertIs<TestException>(uncaughtExceptions.single())
+        assertEquals("x", exception.message)
+    }
+
+    @Test
+    fun testExceptionCaptorCleanedUpOnPreliminaryExit() {
+        RunTestTest().testExceptionCaptorCleanedUpOnPreliminaryExit()
+        assertEquals(2, uncaughtExceptions.size, "Expected 2 uncaught exceptions, but got $uncaughtExceptions")
+        for (exception in uncaughtExceptions) {
+            assertIs<TestException>(exception)
+        }
+        assertEquals("A", uncaughtExceptions[0].message)
+        assertEquals("B", uncaughtExceptions[1].message)
+    }
+}