Allow specifying the timeout for runTest

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

@@ -22,10 +22,10 @@ public final class kotlinx/coroutines/test/TestBuildersKt {
 	public static final fun runTest (Lkotlinx/coroutines/test/TestScope;JLkotlin/jvm/functions/Function2;)V
 	public static synthetic fun runTest$default (Lkotlin/coroutines/CoroutineContext;JLkotlin/jvm/functions/Function2;ILjava/lang/Object;)V
 	public static synthetic fun runTest$default (Lkotlinx/coroutines/test/TestCoroutineScope;JLkotlin/jvm/functions/Function2;ILjava/lang/Object;)V
-	public static synthetic fun runTest$default (Lkotlinx/coroutines/test/TestScope;JLkotlin/jvm/functions/Function2;ILjava/lang/Object;)V
 	public static final fun runTest-8Mi8wO0 (Lkotlin/coroutines/CoroutineContext;JLkotlin/jvm/functions/Function2;)V
 	public static final fun runTest-8Mi8wO0 (Lkotlinx/coroutines/test/TestScope;JLkotlin/jvm/functions/Function2;)V
 	public static synthetic fun runTest-8Mi8wO0$default (Lkotlin/coroutines/CoroutineContext;JLkotlin/jvm/functions/Function2;ILjava/lang/Object;)V
+	public static synthetic fun runTest-8Mi8wO0$default (Lkotlinx/coroutines/test/TestScope;JLkotlin/jvm/functions/Function2;ILjava/lang/Object;)V
 	public static final fun runTestWithLegacyScope (Lkotlin/coroutines/CoroutineContext;JLkotlin/jvm/functions/Function2;)V
 	public static synthetic fun runTestWithLegacyScope$default (Lkotlin/coroutines/CoroutineContext;JLkotlin/jvm/functions/Function2;ILjava/lang/Object;)V
 }
@@ -66,6 +66,7 @@ public final class kotlinx/coroutines/test/TestCoroutineScheduler : kotlin/corou
 	public static final field Key Lkotlinx/coroutines/test/TestCoroutineScheduler$Key;
 	public fun <init> ()V
 	public final fun advanceTimeBy (J)V
+	public final fun advanceTimeBy-LRDsOJo (J)V
 	public final fun advanceUntilIdle ()V
 	public final fun getCurrentTime ()J
 	public final fun getTimeSource ()Lkotlin/time/TimeSource;
@@ -116,6 +117,7 @@ public final class kotlinx/coroutines/test/TestScopeKt {
 	public static final fun TestScope (Lkotlin/coroutines/CoroutineContext;)Lkotlinx/coroutines/test/TestScope;
 	public static synthetic fun TestScope$default (Lkotlin/coroutines/CoroutineContext;ILjava/lang/Object;)Lkotlinx/coroutines/test/TestScope;
 	public static final fun advanceTimeBy (Lkotlinx/coroutines/test/TestScope;J)V
+	public static final fun advanceTimeBy-HG0u8IE (Lkotlinx/coroutines/test/TestScope;J)V
 	public static final fun advanceUntilIdle (Lkotlinx/coroutines/test/TestScope;)V
 	public static final fun getCurrentTime (Lkotlinx/coroutines/test/TestScope;)J
 	public static final fun getTestTimeSource (Lkotlinx/coroutines/test/TestScope;)Lkotlin/time/TimeSource;

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

@@ -7,11 +7,14 @@
 package kotlinx.coroutines.test
 
 import kotlinx.coroutines.*
+import kotlinx.coroutines.flow.*
 import kotlinx.coroutines.selects.*
 import kotlin.coroutines.*
 import kotlin.jvm.*
 import kotlin.time.*
 import kotlin.time.Duration.Companion.milliseconds
+import kotlin.time.Duration.Companion.seconds
+import kotlinx.coroutines.internal.*
 
 /**
  * A test result.
@@ -118,6 +121,18 @@ public expect class TestResult
  *
  * If the created coroutine completes with an exception, then this exception will be thrown at the end of the test.
  *
+ * #### Timing out
+ *
+ * There's a built-in timeout of 10 seconds for the test body. If the test body doesn't complete within this time,
+ * then the test fails with an [AssertionError]. The timeout can be changed by setting the [timeout] parameter.
+ *
+ * The test finishes by the timeout procedure cancelling the test body. If the code inside the test body does not
+ * respond to cancellation, we will not be able to make the test execution stop, in which case, the test will hang
+ * despite our best efforts to terminate it.
+ *
+ * On the JVM, if `DebugProbes` from the `kotlinx-coroutines-debug` module are installed, the current dump of the
+ * coroutines' stack is printed to the console on timeout.
+ *
  * #### Reported exceptions
  *
  * Unhandled exceptions will be thrown at the end of the test.
@@ -131,12 +146,6 @@ public expect class TestResult
  * 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).
  *
- * In the general case, if there are active jobs, it's impossible to detect if they are going to complete eventually due
- * to the asynchronous nature of coroutines. In order to prevent tests hanging in this scenario, [runTest] will wait
- * for [dispatchTimeout] (by default, 60 seconds) from the moment when [TestCoroutineScheduler] becomes
- * idle before throwing [AssertionError]. If some dispatcher linked to [TestCoroutineScheduler] receives a
- * task during that time, the timer gets reset.
- *
  * ### Configuration
  *
  * [context] can be used to affect the environment of the code under test. Beside just being passed to the coroutine
@@ -148,12 +157,12 @@ public expect class TestResult
 @ExperimentalCoroutinesApi
 public fun runTest(
     context: CoroutineContext = EmptyCoroutineContext,
-    dispatchTimeout: Duration = DEFAULT_DISPATCH_TIMEOUT,
+    timeout: Duration = DEFAULT_TIMEOUT,
     testBody: suspend TestScope.() -> Unit
 ): TestResult {
     if (context[RunningInRunTest] != null)
         throw IllegalStateException("Calls to `runTest` can't be nested. Please read the docs on `TestResult` for details.")
-    return TestScope(context + RunningInRunTest).runTest(dispatchTimeout, testBody)
+    return TestScope(context + RunningInRunTest).runTest(timeout, testBody)
 }
 
 /**
@@ -269,46 +278,126 @@ public fun runTest(
  *
  * @throws IllegalArgumentException if the [context] is invalid. See the [TestScope] constructor docs for details.
  */
-@ExperimentalCoroutinesApi
+@Deprecated(
+    "Define a total timeout for the whole test instead of using dispatchTimeoutMs. " +
+        "Warning: the proposed replacement is not identical as it uses 'dispatchTimeoutMs' as the timeout for the whole test!",
+    ReplaceWith("runTest(context, timeout = dispatchTimeoutMs.milliseconds, testBody)",
+        "kotlin.time.Duration.Companion.milliseconds"),
+    DeprecationLevel.WARNING
+)
 public fun runTest(
     context: CoroutineContext = EmptyCoroutineContext,
     dispatchTimeoutMs: Long,
     testBody: suspend TestScope.() -> Unit
-): TestResult = runTest(
-    context = context,
-    dispatchTimeout = dispatchTimeoutMs.milliseconds,
-    testBody = testBody
-)
+): TestResult {
+    if (context[RunningInRunTest] != null)
+        throw IllegalStateException("Calls to `runTest` can't be nested. Please read the docs on `TestResult` for details.")
+    @Suppress("DEPRECATION")
+    return TestScope(context + RunningInRunTest).runTest(dispatchTimeoutMs = dispatchTimeoutMs, testBody)
+}
 
 /**
- * Performs [runTest] on an existing [TestScope].
+ * Performs [runTest] on an existing [TestScope]. See the documentation for [runTest] for details.
  */
 @ExperimentalCoroutinesApi
 public fun TestScope.runTest(
-    dispatchTimeout: Duration,
+    timeout: Duration = DEFAULT_TIMEOUT,
     testBody: suspend TestScope.() -> Unit
 ): TestResult = asSpecificImplementation().let {
     it.enter()
     createTestResult {
-        runTestCoroutine(it, dispatchTimeout, TestScopeImpl::tryGetCompletionCause, testBody) {
+        /** TODO: moving this [AbstractCoroutine.start] call outside [createTestResult] fails on JS. */
+        it.start(CoroutineStart.UNDISPATCHED, it) {
+            testBody()
+        }
+        /**
+         * We run the tasks in the test coroutine using [Dispatchers.Default]. On JS, this does nothing particularly,
+         * but on the JVM and Native, this means that the timeout can be processed even while the test runner is busy
+         * doing some synchronous work.
+         */
+        val workRunner = launch(Dispatchers.Default + CoroutineName("kotlinx.coroutines.test runner")) {
+            while (true) {
+                val executedSomething = testScheduler.tryRunNextTaskUnless { !isActive }
+                if (executedSomething) {
+                    /** yield to check for cancellation. On JS, we can't use [ensureActive] here, as the cancellation
+                     * procedure needs a chance to run concurrently. */
+                    yield()
+                } else {
+                    // no more tasks, we should suspend until there are some more
+                    testScheduler.receiveDispatchEvent()
+                }
+            }
+        }
+        var timeoutError: Throwable? = null
+        try {
+            withTimeout(timeout) {
+                it.join()
+                workRunner.cancelAndJoin()
+            }
+        } catch (_: TimeoutCancellationException) {
+            val activeChildren = it.children.filter(Job::isActive).toList()
+            val completionCause = if (it.isCancelled) it.tryGetCompletionCause() else null
+            var message = "After waiting for $timeout"
+            if (completionCause == null)
+                message += ", the test coroutine is not completing"
+            if (activeChildren.isNotEmpty())
+                message += ", there were active child jobs: $activeChildren"
+            if (completionCause != null && activeChildren.isEmpty()) {
+                message += if (it.isCompleted)
+                    ", the test coroutine completed"
+                else
+                    ", the test coroutine was not completed"
+            }
+            timeoutError = UncompletedCoroutinesError(message)
+            dumpCoroutines()
+            val cancellationException = CancellationException("The test timed out")
+            (it as Job).cancel(cancellationException)
+            // we can't abandon the work we're doing, so if it hanged, we'll still hang, despite the timeout.
+            it.join()
+            val completion = it.getCompletionExceptionOrNull()
+            if (completion != null && completion !== cancellationException) {
+                timeoutError.addSuppressed(completion)
+            }
+            workRunner.cancelAndJoin()
+        } finally {
             backgroundScope.cancel()
             testScheduler.advanceUntilIdleOr { false }
-            it.leave()
+            val uncaughtExceptions = it.leave()
+            throwAll(timeoutError ?: it.getCompletionExceptionOrNull(), uncaughtExceptions)
         }
     }
 }
 
 /**
  * Performs [runTest] on an existing [TestScope].
+ *
+ * In the general case, if there are active jobs, it's impossible to detect if they are going to complete eventually due
+ * to the asynchronous nature of coroutines. In order to prevent tests hanging in this scenario, [runTest] will wait
+ * for [dispatchTimeoutMs] from the moment when [TestCoroutineScheduler] becomes
+ * idle before throwing [AssertionError]. If some dispatcher linked to [TestCoroutineScheduler] receives a
+ * task during that time, the timer gets reset.
  */
-@ExperimentalCoroutinesApi
+@Deprecated(
+    "Define a total timeout for the whole test instead of using dispatchTimeoutMs. " +
+        "Warning: the proposed replacement is not identical as it uses 'dispatchTimeoutMs' as the timeout for the whole test!",
+    ReplaceWith("this.runTest(timeout = dispatchTimeoutMs.milliseconds, testBody)",
+        "kotlin.time.Duration.Companion.milliseconds"),
+    DeprecationLevel.WARNING
+)
 public fun TestScope.runTest(
-    dispatchTimeoutMs: Long = DEFAULT_DISPATCH_TIMEOUT_MS,
+    dispatchTimeoutMs: Long,
     testBody: suspend TestScope.() -> Unit
-): TestResult = runTest(
-    dispatchTimeout = dispatchTimeoutMs.milliseconds,
-    testBody = testBody
-)
+): TestResult = asSpecificImplementation().let {
+    it.enter()
+    @Suppress("DEPRECATION")
+    createTestResult {
+        runTestCoroutineLegacy(it, dispatchTimeoutMs.milliseconds, TestScopeImpl::tryGetCompletionCause, testBody) {
+            backgroundScope.cancel()
+            testScheduler.advanceUntilIdleOr { false }
+            it.legacyLeave()
+        }
+    }
+}
 
 /**
  * Runs [testProcedure], creating a [TestResult].
@@ -327,18 +416,23 @@ internal object RunningInRunTest : CoroutineContext.Key<RunningInRunTest>, Corou
 /** The default timeout to use when waiting for asynchronous completions of the coroutines managed by
  * a [TestCoroutineScheduler]. */
 internal const val DEFAULT_DISPATCH_TIMEOUT_MS = 60_000L
-internal val DEFAULT_DISPATCH_TIMEOUT = DEFAULT_DISPATCH_TIMEOUT_MS.milliseconds
+
+/**
+ * The default timeout to use when running a test.
+ */
+internal val DEFAULT_TIMEOUT = 10.seconds
 
 /**
  * Run the [body][testBody] of the [test coroutine][coroutine], waiting for asynchronous completions for at most
- * [dispatchTimeout], and performing the [cleanup] procedure at the end.
+ * [dispatchTimeout] and performing the [cleanup] procedure at the end.
  *
  * [tryGetCompletionCause] is the [JobSupport.completionCause], which is passed explicitly because it is protected.
  *
  * The [cleanup] procedure may either throw [UncompletedCoroutinesError] to denote that child coroutines were leaked, or
  * return a list of uncaught exceptions that should be reported at the end of the test.
  */
-internal suspend fun <T: AbstractCoroutine<Unit>> CoroutineScope.runTestCoroutine(
+@Deprecated("Used for support of legacy behavior")
+internal suspend fun <T : AbstractCoroutine<Unit>> CoroutineScope.runTestCoroutineLegacy(
     coroutine: T,
     dispatchTimeout: Duration,
     tryGetCompletionCause: T.() -> Throwable?,
@@ -351,6 +445,8 @@ internal suspend fun <T: AbstractCoroutine<Unit>> CoroutineScope.runTestCoroutin
         testBody()
     }
     /**
+     * This is the legacy behavior, kept for now for compatibility only.
+     *
      * The general procedure here is as follows:
      * 1. Try running the work that the scheduler knows about, both background and foreground.
      *
@@ -376,16 +472,22 @@ internal suspend fun <T: AbstractCoroutine<Unit>> CoroutineScope.runTestCoroutin
         scheduler.advanceUntilIdle()
         if (coroutine.isCompleted) {
             /* don't even enter `withTimeout`; this allows to use a timeout of zero to check that there are no
-               non-trivial dispatches. */
+           non-trivial dispatches. */
             completed = true
             continue
         }
         // in case progress depends on some background work, we need to keep spinning it.
         val backgroundWorkRunner = launch(CoroutineName("background work runner")) {
             while (true) {
-                scheduler.tryRunNextTaskUnless { !isActive }
-                // yield so that the `select` below has a chance to check if its conditions are fulfilled
-                yield()
+                val executedSomething = scheduler.tryRunNextTaskUnless { !isActive }
+                if (executedSomething) {
+                    // yield so that the `select` below has a chance to finish successfully or time out
+                    yield()
+                } else {
+                    // no more tasks, we should suspend until there are some more.
+                    // this doesn't interfere with the `select` below, because different channels are used.
+                    scheduler.receiveDispatchEvent()
+                }
             }
         }
         try {
@@ -394,11 +496,11 @@ internal suspend fun <T: AbstractCoroutine<Unit>> CoroutineScope.runTestCoroutin
                     // observe that someone completed the test coroutine and leave without waiting for the timeout
                     completed = true
                 }
-                scheduler.onDispatchEvent {
+                scheduler.onDispatchEventForeground {
                     // we received knowledge that `scheduler` observed a dispatch event, so we reset the timeout
                 }
                 onTimeout(dispatchTimeout) {
-                    handleTimeout(coroutine, dispatchTimeout, tryGetCompletionCause, cleanup)
+                    throw handleTimeout(coroutine, dispatchTimeout, tryGetCompletionCause, cleanup)
                 }
             }
         } finally {
@@ -412,21 +514,20 @@ internal suspend fun <T: AbstractCoroutine<Unit>> CoroutineScope.runTestCoroutin
             // it's normal that some jobs are not completed if the test body has failed, won't clutter the output
             emptyList()
         }
-        (listOf(exception) + exceptions).throwAll()
+        throwAll(exception, exceptions)
     }
-    cleanup().throwAll()
+    throwAll(null, cleanup())
 }
 
 /**
- * Invoked on timeout in [runTest]. Almost always just builds a nice [UncompletedCoroutinesError] and throws it.
- * However, sometimes it detects that the coroutine completed, in which case it returns normally.
+ * Invoked on timeout in [runTest]. Just builds a nice [UncompletedCoroutinesError] and returns it.
  */
-private inline fun<T: AbstractCoroutine<Unit>> handleTimeout(
+private inline fun <T : AbstractCoroutine<Unit>> handleTimeout(
     coroutine: T,
     dispatchTimeout: Duration,
     tryGetCompletionCause: T.() -> Throwable?,
     cleanup: () -> List<Throwable>,
-) {
+): AssertionError {
     val uncaughtExceptions = try {
         cleanup()
     } catch (e: UncompletedCoroutinesError) {
@@ -441,20 +542,29 @@ private inline fun<T: AbstractCoroutine<Unit>> handleTimeout(
     if (activeChildren.isNotEmpty())
         message += ", there were active child jobs: $activeChildren"
     if (completionCause != null && activeChildren.isEmpty()) {
-        if (coroutine.isCompleted)
-            return
-        // TODO: can this really ever happen?
-        message += ", the test coroutine was not completed"
+        message += if (coroutine.isCompleted)
+            ", the test coroutine completed"
+        else
+            ", the test coroutine was not completed"
     }
     val error = UncompletedCoroutinesError(message)
     completionCause?.let { cause -> error.addSuppressed(cause) }
     uncaughtExceptions.forEach { error.addSuppressed(it) }
-    throw error
+    return error
 }
 
-internal fun List<Throwable>.throwAll() {
-    firstOrNull()?.apply {
-        drop(1).forEach { addSuppressed(it) }
-        throw this
+internal fun throwAll(head: Throwable?, other: List<Throwable>) {
+    if (head != null) {
+        other.forEach { head.addSuppressed(it) }
+        throw head
+    } else {
+        with(other) {
+            firstOrNull()?.apply {
+                drop(1).forEach { addSuppressed(it) }
+                throw this
+            }
+        }
     }
 }
+
+internal expect fun dumpCoroutines()