Change exception handling in the test module

Author: dkhalanskyjb

kotlinx-coroutines-test/api/kotlinx-coroutines-test.api

@@ -51,7 +51,7 @@ public final class kotlinx/coroutines/test/TestCoroutineDispatcher : kotlinx/cor
 
 public final class kotlinx/coroutines/test/TestCoroutineExceptionHandler : kotlin/coroutines/AbstractCoroutineContextElement, kotlinx/coroutines/CoroutineExceptionHandler, kotlinx/coroutines/test/UncaughtExceptionCaptor {
 	public fun <init> ()V
-	public fun cleanupTestCoroutinesCaptor ()V
+	public fun cleanupTestCoroutines ()V
 	public fun getUncaughtExceptions ()Ljava/util/List;
 	public fun handleException (Lkotlin/coroutines/CoroutineContext;Ljava/lang/Throwable;)V
 }
@@ -68,9 +68,10 @@ public final class kotlinx/coroutines/test/TestCoroutineScheduler : kotlin/corou
 public final class kotlinx/coroutines/test/TestCoroutineScheduler$Key : kotlin/coroutines/CoroutineContext$Key {
 }
 
-public abstract interface class kotlinx/coroutines/test/TestCoroutineScope : kotlinx/coroutines/CoroutineScope, kotlinx/coroutines/test/UncaughtExceptionCaptor {
+public abstract interface class kotlinx/coroutines/test/TestCoroutineScope : kotlinx/coroutines/CoroutineScope {
 	public abstract fun cleanupTestCoroutines ()V
 	public abstract fun getTestScheduler ()Lkotlinx/coroutines/test/TestCoroutineScheduler;
+	public abstract fun reportException (Ljava/lang/Throwable;)V
 }
 
 public final class kotlinx/coroutines/test/TestCoroutineScopeKt {
@@ -81,6 +82,7 @@ public final class kotlinx/coroutines/test/TestCoroutineScopeKt {
 	public static final fun createTestCoroutineScope (Lkotlin/coroutines/CoroutineContext;)Lkotlinx/coroutines/test/TestCoroutineScope;
 	public static synthetic fun createTestCoroutineScope$default (Lkotlin/coroutines/CoroutineContext;ILjava/lang/Object;)Lkotlinx/coroutines/test/TestCoroutineScope;
 	public static final fun getCurrentTime (Lkotlinx/coroutines/test/TestCoroutineScope;)J
+	public static final fun getUncaughtExceptions (Lkotlinx/coroutines/test/TestCoroutineScope;)Ljava/util/List;
 	public static final fun pauseDispatcher (Lkotlinx/coroutines/test/TestCoroutineScope;)V
 	public static final fun pauseDispatcher (Lkotlinx/coroutines/test/TestCoroutineScope;Lkotlin/jvm/functions/Function1;Lkotlin/coroutines/Continuation;)Ljava/lang/Object;
 	public static final fun resumeDispatcher (Lkotlinx/coroutines/test/TestCoroutineScope;)V
@@ -100,8 +102,13 @@ public final class kotlinx/coroutines/test/TestDispatchers {
 	public static final fun setMain (Lkotlinx/coroutines/Dispatchers;Lkotlinx/coroutines/CoroutineDispatcher;)V
 }
 
+public final class kotlinx/coroutines/test/TestExceptionHandlerKt {
+	public static final fun TestExceptionHandler (Lkotlinx/coroutines/test/TestCoroutineScope;Lkotlin/jvm/functions/Function3;)Lkotlinx/coroutines/CoroutineExceptionHandler;
+	public static synthetic fun TestExceptionHandler$default (Lkotlinx/coroutines/test/TestCoroutineScope;Lkotlin/jvm/functions/Function3;ILjava/lang/Object;)Lkotlinx/coroutines/CoroutineExceptionHandler;
+}
+
 public abstract interface class kotlinx/coroutines/test/UncaughtExceptionCaptor {
-	public abstract fun cleanupTestCoroutinesCaptor ()V
+	public abstract fun cleanupTestCoroutines ()V
 	public abstract fun getUncaughtExceptions ()Ljava/util/List;
 }
 

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

@@ -280,11 +280,11 @@ private const val DEFAULT_DISPATCH_TIMEOUT_MS = 60_000L
 
 private class TestBodyCoroutine<T>(
     private val testScope: TestCoroutineScope,
-) : AbstractCoroutine<T>(testScope.coroutineContext, initParentJob = true, active = true), TestCoroutineScope,
-    UncaughtExceptionCaptor by testScope.coroutineContext.uncaughtExceptionCaptor
+) : AbstractCoroutine<T>(testScope.coroutineContext, initParentJob = true, active = true), TestCoroutineScope
 {
     override val testScheduler get() = testScope.testScheduler
 
     override fun cleanupTestCoroutines() = testScope.cleanupTestCoroutines()
 
+    override fun reportException(throwable: Throwable) = testScope.reportException(throwable)
 }

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

@@ -11,13 +11,16 @@ import kotlin.coroutines.*
 /**
  * Access uncaught coroutine exceptions captured during test execution.
  */
-@ExperimentalCoroutinesApi
+@Deprecated(
+    "Consider whether a `TestExceptionHandler` would work instead. If not, please report your use case at https://github.com/Kotlin/kotlinx.coroutines/issues.",
+    level = DeprecationLevel.WARNING
+)
 public interface UncaughtExceptionCaptor {
     /**
      * List of uncaught coroutine exceptions.
      *
      * The returned list is a copy of the currently caught exceptions.
-     * During [cleanupTestCoroutinesCaptor] the first element of this list is rethrown if it is not empty.
+     * During [cleanupTestCoroutines] the first element of this list is rethrown if it is not empty.
      */
     public val uncaughtExceptions: List<Throwable>
 
@@ -29,33 +32,39 @@ public interface UncaughtExceptionCaptor {
      *
      * @throws Throwable the first uncaught exception, if there are any uncaught exceptions.
      */
-    public fun cleanupTestCoroutinesCaptor()
+    public fun cleanupTestCoroutines()
 }
 
 /**
  * An exception handler that captures uncaught exceptions in tests.
  */
-@ExperimentalCoroutinesApi
+@Deprecated("You can use `TestExceptionHandler` to define an exception handler that is linked to " +
+    "a `TestCoroutineScope`, or define your own `CoroutineExceptionHandler` if you just need to handle uncaught " +
+    "exceptions without a special `TestCoroutineScope` integration.", level = DeprecationLevel.WARNING)
 public class TestCoroutineExceptionHandler :
-    AbstractCoroutineContextElement(CoroutineExceptionHandler), UncaughtExceptionCaptor, CoroutineExceptionHandler
+    AbstractCoroutineContextElement(CoroutineExceptionHandler), CoroutineExceptionHandler, UncaughtExceptionCaptor
 {
     private val _exceptions = mutableListOf<Throwable>()
     private val _lock = SynchronizedObject()
+    private var _coroutinesCleanedUp = false
 
-    /** @suppress **/
+    @Suppress("INVISIBLE_MEMBER")
     override fun handleException(context: CoroutineContext, exception: Throwable) {
         synchronized(_lock) {
+            if (_coroutinesCleanedUp) {
+                handleCoroutineExceptionImpl(context, exception)
+                return
+            }
             _exceptions += exception
         }
     }
 
-    /** @suppress **/
-    override val uncaughtExceptions: List<Throwable>
+    public override val uncaughtExceptions: List<Throwable>
         get() = synchronized(_lock) { _exceptions.toList() }
 
-    /** @suppress **/
-    override fun cleanupTestCoroutinesCaptor() {
+    public override fun cleanupTestCoroutines() {
         synchronized(_lock) {
+            _coroutinesCleanedUp = true
             val exception = _exceptions.firstOrNull() ?: return
             // log the rest
             _exceptions.drop(1).forEach { it.printStackTrace() }

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

@@ -5,21 +5,33 @@
 package kotlinx.coroutines.test
 
 import kotlinx.coroutines.*
+import kotlinx.coroutines.internal.*
 import kotlin.coroutines.*
 
 /**
  * A scope which provides detailed control over the execution of coroutines for tests.
  */
 @ExperimentalCoroutinesApi
-public sealed interface TestCoroutineScope: CoroutineScope, UncaughtExceptionCaptor {
+public sealed interface TestCoroutineScope: CoroutineScope {
     /**
      * Called after the test completes.
      *
-     * Calls [UncaughtExceptionCaptor.cleanupTestCoroutinesCaptor] and [DelayController.cleanupTestCoroutines].
-     * If a new job was created for this scope, the job is completed.
+     * * It checks that there were no uncaught exceptions reported via [reportException]. If there were any, then the
+     *   first one is thrown, whereas the rest are printed to the standard output or the standard error output
+     *   (see [Throwable.printStackTrace]).
+     * * It runs the tasks pending in the scheduler at the current time. If there are any uncompleted tasks afterwards,
+     *   it fails with [UncompletedCoroutinesError].
+     * * It checks whether some new child [Job]s were created but not completed since this [TestCoroutineScope] was
+     *   created. If so, it fails with [UncompletedCoroutinesError].
+     *
+     * For backwards compatibility, if the [CoroutineExceptionHandler] is an [UncaughtExceptionCaptor], its
+     * [TestCoroutineExceptionHandler.cleanupTestCoroutines] behavior is performed.
+     * Likewise, if the [ContinuationInterceptor] is a [DelayController], its [DelayController.cleanupTestCoroutines]
+     * is called.
      *
      * @throws Throwable the first uncaught exception, if there are any uncaught exceptions.
      * @throws AssertionError if any pending tasks are active.
+     * @throws IllegalStateException if called more than once.
      */
     @ExperimentalCoroutinesApi
     public fun cleanupTestCoroutines()
@@ -29,14 +41,36 @@ public sealed interface TestCoroutineScope: CoroutineScope, UncaughtExceptionCap
      */
     @ExperimentalCoroutinesApi
     public val testScheduler: TestCoroutineScheduler
+
+    /**
+     * Reports an exception so that it is thrown on [cleanupTestCoroutines].
+     *
+     * If several exceptions are reported, only the first one will be thrown, and the other ones will be suppressed by
+     * it.
+     *
+     * @throws IllegalStateException with the [Throwable.cause] set to [throwable] if [cleanupTestCoroutines] was
+     * already called.
+     */
+    @ExperimentalCoroutinesApi
+    public fun reportException(throwable: Throwable)
 }
 
 private class TestCoroutineScopeImpl(
     override val coroutineContext: CoroutineContext
-):
-    TestCoroutineScope,
-    UncaughtExceptionCaptor by coroutineContext.uncaughtExceptionCaptor
+): TestCoroutineScope
 {
+    private val lock = SynchronizedObject()
+    private var exceptions = mutableListOf<Throwable>()
+    private var cleanedUp = false
+
+    override fun reportException(throwable: Throwable) {
+        synchronized(lock) {
+            if (cleanedUp)
+                throw ExceptionReportAfterCleanup(throwable)
+            exceptions.add(throwable)
+        }
+    }
+
     override val testScheduler: TestCoroutineScheduler
         get() = coroutineContext[TestCoroutineScheduler]!!
 
@@ -56,7 +90,16 @@ private class TestCoroutineScopeImpl(
             testScheduler.runCurrent()
             !testScheduler.isIdle()
         }
-        coroutineContext.uncaughtExceptionCaptor.cleanupTestCoroutinesCaptor()
+        (coroutineContext[CoroutineExceptionHandler] as? UncaughtExceptionCaptor)?.cleanupTestCoroutines()
+        synchronized(lock) {
+            if (cleanedUp)
+                throw IllegalStateException("Attempting to clean up a test coroutine scope more than once.")
+            cleanedUp = true
+        }
+        exceptions.firstOrNull()?.let { toThrow ->
+            exceptions.drop(1).forEach { toThrow.addSuppressed(it) }
+            throw toThrow
+        }
         if (hasUnfinishedJobs)
             throw UncompletedCoroutinesError(
                 "Unfinished coroutines during teardown. Ensure all coroutines are" +
@@ -68,6 +111,11 @@ private class TestCoroutineScopeImpl(
     }
 }
 
+internal class ExceptionReportAfterCleanup(cause: Throwable): IllegalStateException(
+    "Attempting to report an uncaught exception after the test coroutine scope was already cleaned up",
+    cause
+)
+
 private fun CoroutineContext.activeJobs(): Set<Job> {
     return checkNotNull(this[Job]).children.filter { it.isActive }.toSet()
 }
@@ -79,13 +127,13 @@ private fun CoroutineContext.activeJobs(): Set<Job> {
  */
 @Deprecated("This constructs a `TestCoroutineScope` with a deprecated `CoroutineDispatcher` by default. " +
     "Please use `createTestCoroutineScope` instead.",
-    ReplaceWith("createTestCoroutineScope(TestCoroutineDispatcher() + context)",
+    ReplaceWith("createTestCoroutineScope(TestCoroutineDispatcher() + TestCoroutineExceptionHandler() + context)",
         "kotlin.coroutines.EmptyCoroutineContext"),
     level = DeprecationLevel.WARNING
 )
 public fun TestCoroutineScope(context: CoroutineContext = EmptyCoroutineContext): TestCoroutineScope {
     val scheduler = context[TestCoroutineScheduler] ?: TestCoroutineScheduler()
-    return createTestCoroutineScope(TestCoroutineDispatcher(scheduler) + context)
+    return createTestCoroutineScope(TestCoroutineDispatcher(scheduler) + TestCoroutineExceptionHandler() + context)
 }
 
 /**
@@ -126,26 +174,28 @@ public fun createTestCoroutineScope(context: CoroutineContext = EmptyCoroutineCo
         }
         else -> throw IllegalArgumentException("Dispatcher must implement TestDispatcher: $dispatcher")
     }
-    val exceptionHandler = context[CoroutineExceptionHandler].run {
-        this?.let {
-            require(this is UncaughtExceptionCaptor) {
-                "coroutineExceptionHandler must implement UncaughtExceptionCaptor: $context"
-            }
+    val linkedHandler: TestExceptionHandlerContextElement?
+    val exceptionHandler: CoroutineExceptionHandler
+    val handlerOwner = Any()
+    when (val exceptionHandlerInCtx = context[CoroutineExceptionHandler]) {
+        null -> {
+            linkedHandler = TestExceptionHandlerContextElement(
+                { _, throwable -> reportException(throwable) },
+                null,
+                handlerOwner)
+            exceptionHandler = linkedHandler
+        }
+        else -> {
+            linkedHandler = (exceptionHandlerInCtx as? TestExceptionHandlerContextElement
+                )?.claimOwnershipOrCopy(handlerOwner)
+            exceptionHandler = linkedHandler ?: exceptionHandlerInCtx
         }
-        this ?: TestCoroutineExceptionHandler()
     }
     val job: Job = context[Job] ?: SupervisorJob()
     return TestCoroutineScopeImpl(context + scheduler + dispatcher + exceptionHandler + job)
+        .also { linkedHandler?.registerTestCoroutineScope(handlerOwner, it) }
 }
 
-internal inline val CoroutineContext.uncaughtExceptionCaptor: UncaughtExceptionCaptor
-    get() {
-        val handler = this[CoroutineExceptionHandler]
-        return handler as? UncaughtExceptionCaptor ?: throw IllegalArgumentException(
-            "TestCoroutineScope requires a UncaughtExceptionCaptor such as " +
-                "TestCoroutineExceptionHandler as the CoroutineExceptionHandler"
-        )
-    }
 
 private inline val CoroutineContext.delayController: DelayController?
     get() {
@@ -238,6 +288,24 @@ public fun TestCoroutineScope.resumeDispatcher() {
     delayControllerForPausing.resumeDispatcher()
 }
 
+/**
+ * List of uncaught coroutine exceptions, for backward compatibility.
+ *
+ * The returned list is a copy of the exceptions caught during execution.
+ * During [TestCoroutineScope.cleanupTestCoroutines] the first element of this list is rethrown if it is not empty.
+ *
+ * Exceptions are only collected in this list if the [UncaughtExceptionCaptor] is in the test context.
+ */
+@Deprecated(
+    "This list is only populated if `UncaughtExceptionCaptor` is in the test context, and so can be " +
+        "easily misused. It is only present for backward compatibility and will be removed in the subsequent " +
+        "releases. If you need to check the list of exceptions, please consider creating your own " +
+        "`CoroutineExceptionHandler`.",
+    level = DeprecationLevel.WARNING)
+public val TestCoroutineScope.uncaughtExceptions: List<Throwable>
+    get() = (coroutineContext[CoroutineExceptionHandler] as? UncaughtExceptionCaptor)?.uncaughtExceptions
+        ?: emptyList()
+
 private val TestCoroutineScope.delayControllerForPausing: DelayController
     get() = coroutineContext.delayController
         ?: throw IllegalStateException("This scope isn't able to pause its dispatchers")