Allow specifying the timeout for runTest (#3603)

Deprecate `dispatchTimeoutMs`, as this is a confusing
implementation detail that made it to the final API.

We use the fact that the `runTest(Duration)` overload was never
published, so we can reuse it to have the `Duration` mean the
whole-test timeout in a backward-compatible manner.

Author: dkhalanskyjb

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;
@@ -117,6 +118,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/build.gradle.kts

@@ -1,9 +1,9 @@
-import org.jetbrains.kotlin.gradle.plugin.mpp.*
-
 /*
  * Copyright 2016-2021 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license.
  */
 
+import org.jetbrains.kotlin.gradle.plugin.mpp.*
+
 val experimentalAnnotations = listOf(
     "kotlin.Experimental",
     "kotlinx.coroutines.ExperimentalCoroutinesApi",
@@ -19,4 +19,12 @@ kotlin {
             binaryOptions["memoryModel"] = "experimental"
         }
     }
+
+    sourceSets {
+        jvmTest {
+            dependencies {
+                implementation(project(":kotlinx-coroutines-debug"))
+            }
+        }
+    }
 }

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

@@ -13,6 +13,7 @@ import kotlinx.coroutines.selects.*
 import kotlin.coroutines.*
 import kotlin.jvm.*
 import kotlin.time.*
+import kotlin.time.Duration.Companion.milliseconds
 
 /**
  * This is a scheduler for coroutines used in tests, providing the delay-skipping behavior.
@@ -49,6 +50,9 @@ public class TestCoroutineScheduler : AbstractCoroutineContextElement(TestCorout
         get() = synchronized(lock) { field }
         private set
 
+    /** A channel for notifying about the fact that a foreground work dispatch recently happened. */
+    private val dispatchEventsForeground: Channel<Unit> = Channel(CONFLATED)
+
     /** A channel for notifying about the fact that a dispatch recently happened. */
     private val dispatchEvents: Channel<Unit> = Channel(CONFLATED)
 
@@ -73,8 +77,8 @@ public class TestCoroutineScheduler : AbstractCoroutineContextElement(TestCorout
             val time = addClamping(currentTime, timeDeltaMillis)
             val event = TestDispatchEvent(dispatcher, count, time, marker as Any, isForeground) { isCancelled(marker) }
             events.addLast(event)
-            /** can't be moved above: otherwise, [onDispatchEvent] could consume the token sent here before there's
-             * actually anything in the event queue. */
+            /** can't be moved above: otherwise, [onDispatchEventForeground] or [onDispatchEvent] could consume the
+             * token sent here before there's actually anything in the event queue. */
             sendDispatchEvent(context)
             DisposableHandle {
                 synchronized(lock) {
@@ -150,13 +154,22 @@ public class TestCoroutineScheduler : AbstractCoroutineContextElement(TestCorout
      * * Overflowing the target time used to lead to nothing being done, but will now run the tasks scheduled at up to
      *   (but not including) [Long.MAX_VALUE].
      *
-     * @throws IllegalStateException if passed a negative [delay][delayTimeMillis].
+     * @throws IllegalArgumentException if passed a negative [delay][delayTimeMillis].
+     */
+    @ExperimentalCoroutinesApi
+    public fun advanceTimeBy(delayTimeMillis: Long): Unit = advanceTimeBy(delayTimeMillis.milliseconds)
+
+    /**
+     * Moves the virtual clock of this dispatcher forward by [the specified amount][delayTime], running the
+     * scheduled tasks in the meantime.
+     *
+     * @throws IllegalArgumentException if passed a negative [delay][delayTime].
      */
     @ExperimentalCoroutinesApi
-    public fun advanceTimeBy(delayTimeMillis: Long) {
-        require(delayTimeMillis >= 0) { "Can not advance time by a negative delay: $delayTimeMillis" }
+    public fun advanceTimeBy(delayTime: Duration) {
+        require(!delayTime.isNegative()) { "Can not advance time by a negative delay: $delayTime" }
         val startingTime = currentTime
-        val targetTime = addClamping(startingTime, delayTimeMillis)
+        val targetTime = addClamping(startingTime, delayTime.inWholeMilliseconds)
         while (true) {
             val event = synchronized(lock) {
                 val timeMark = currentTime
@@ -191,15 +204,26 @@ public class TestCoroutineScheduler : AbstractCoroutineContextElement(TestCorout
      * [context] is the context in which the task will be dispatched.
      */
     internal fun sendDispatchEvent(context: CoroutineContext) {
+        dispatchEvents.trySend(Unit)
         if (context[BackgroundWork] !== BackgroundWork)
-            dispatchEvents.trySend(Unit)
+            dispatchEventsForeground.trySend(Unit)
     }
 
+    /**
+     * Waits for a notification about a dispatch event.
+     */
+    internal suspend fun receiveDispatchEvent() = dispatchEvents.receive()
+
     /**
      * Consumes the knowledge that a dispatch event happened recently.
      */
     internal val onDispatchEvent: SelectClause1<Unit> get() = dispatchEvents.onReceive
 
+    /**
+     * Consumes the knowledge that a foreground work dispatch event happened recently.
+     */
+    internal val onDispatchEventForeground: SelectClause1<Unit> get() = dispatchEventsForeground.onReceive
+
     /**
      * Returns the [TimeSource] representation of the virtual time of this scheduler.
      */

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

@@ -52,9 +52,9 @@ public sealed interface TestScope : CoroutineScope {
      * A scope for background work.
      *
      * This scope is automatically cancelled when the test finishes.
-     * Additionally, while the coroutines in this scope are run as usual when
-     * using [advanceTimeBy] and [runCurrent], [advanceUntilIdle] will stop advancing the virtual time
-     * once only the coroutines in this scope are left unprocessed.
+     * The coroutines in this scope are run as usual when using [advanceTimeBy] and [runCurrent].
+     * [advanceUntilIdle], on the other hand, will stop advancing the virtual time once only the coroutines in this
+     * scope are left unprocessed.
      *
      * Failures in coroutines in this scope do not terminate the test.
      * Instead, they are reported at the end of the test.
@@ -123,6 +123,16 @@ public fun TestScope.runCurrent(): Unit = testScheduler.runCurrent()
 @ExperimentalCoroutinesApi
 public fun TestScope.advanceTimeBy(delayTimeMillis: Long): Unit = testScheduler.advanceTimeBy(delayTimeMillis)
 
+/**
+ * Moves the virtual clock of this dispatcher forward by [the specified amount][delayTime], running the
+ * scheduled tasks in the meantime.
+ *
+ * @throws IllegalStateException if passed a negative [delay][delayTime].
+ * @see TestCoroutineScheduler.advanceTimeBy
+ */
+@ExperimentalCoroutinesApi
+public fun TestScope.advanceTimeBy(delayTime: Duration): Unit = testScheduler.advanceTimeBy(delayTime)
+
 /**
  * The [test scheduler][TestScope.testScheduler] as a [TimeSource].
  * @see TestCoroutineScheduler.timeSource
@@ -230,8 +240,15 @@ internal class TestScopeImpl(context: CoroutineContext) :
         }
     }
 
+    /** Called at the end of the test. May only be called once. Returns the list of caught unhandled exceptions. */
+    fun leave(): List<Throwable> = synchronized(lock) {
+        check(entered && !finished)
+        finished = true
+        uncaughtExceptions
+    }
+
     /** Called at the end of the test. May only be called once. */
-    fun leave(): List<Throwable> {
+    fun legacyLeave(): List<Throwable> {
         val exceptions = synchronized(lock) {
             check(entered && !finished)
             finished = true

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

@@ -9,6 +9,8 @@ import kotlinx.coroutines.internal.*
 import kotlinx.coroutines.flow.*
 import kotlin.coroutines.*
 import kotlin.test.*
+import kotlin.time.*
+import kotlin.time.Duration.Companion.milliseconds
 
 class RunTestTest {
 
@@ -52,7 +54,7 @@ class RunTestTest {
 
     /** Tests that even the dispatch timeout of `0` is fine if all the dispatches go through the same scheduler. */
     @Test
-    fun testRunTestWithZeroTimeoutWithControlledDispatches() = runTest(dispatchTimeoutMs = 0) {
+    fun testRunTestWithZeroDispatchTimeoutWithControlledDispatches() = runTest(dispatchTimeoutMs = 0) {
         // below is some arbitrary concurrent code where all dispatches go through the same scheduler.
         launch {
             delay(2000)
@@ -71,8 +73,13 @@ class RunTestTest {
 
     /** Tests that too low of a dispatch timeout causes crashes. */
     @Test
-    fun testRunTestWithSmallTimeout() = testResultMap({ fn ->
-        assertFailsWith<UncompletedCoroutinesError> { fn() }
+    fun testRunTestWithSmallDispatchTimeout() = testResultMap({ fn ->
+        try {
+            fn()
+            fail("shouldn't be reached")
+        } catch (e: Throwable) {
+            assertIs<UncompletedCoroutinesError>(e)
+        }
     }) {
         runTest(dispatchTimeoutMs = 100) {
             withContext(Dispatchers.Default) {
@@ -83,6 +90,48 @@ class RunTestTest {
         }
     }
 
+    /**
+     * Tests that [runTest] times out after the specified time.
+     */
+    @Test
+    fun testRunTestWithSmallTimeout() = testResultMap({ fn ->
+        try {
+            fn()
+            fail("shouldn't be reached")
+        } catch (e: Throwable) {
+            assertIs<UncompletedCoroutinesError>(e)
+        }
+    }) {
+        runTest(timeout = 100.milliseconds) {
+            withContext(Dispatchers.Default) {
+                delay(10000)
+                3
+            }
+            fail("shouldn't be reached")
+        }
+    }
+
+    /** Tests that [runTest] times out after the specified time, even if the test framework always knows the test is
+     * still doing something. */
+    @Test
+    fun testRunTestWithSmallTimeoutAndManyDispatches() = testResultMap({ fn ->
+        try {
+            fn()
+            fail("shouldn't be reached")
+        } catch (e: Throwable) {
+            assertIs<UncompletedCoroutinesError>(e)
+        }
+    }) {
+        runTest(timeout = 100.milliseconds) {
+            while (true) {
+                withContext(Dispatchers.Default) {
+                    delay(10)
+                    3
+                }
+            }
+        }
+    }
+
     /** Tests that, on timeout, the names of the active coroutines are listed,
      * whereas the names of the completed ones are not. */
     @Test
@@ -119,26 +168,33 @@ class RunTestTest {
         } catch (e: UncompletedCoroutinesError) {
             @Suppress("INVISIBLE_MEMBER")
             val suppressed = unwrap(e).suppressedExceptions
-            assertEquals(1, suppressed.size)
+            assertEquals(1, suppressed.size, "$suppressed")
             assertIs<TestException>(suppressed[0]).also {
                 assertEquals("A", it.message)
             }
         }
     }) {
-        runTest(dispatchTimeoutMs = 10) {
-            launch {
-                withContext(NonCancellable) {
-                    awaitCancellation()
+        runTest(timeout = 10.milliseconds) {
+            launch(start = CoroutineStart.UNDISPATCHED) {
+                withContext(NonCancellable + Dispatchers.Default) {
+                    delay(100.milliseconds)
                 }
             }
-            yield()
             throw TestException("A")
         }
     }
 
     /** Tests that real delays can be accounted for with a large enough dispatch timeout. */
     @Test
-    fun testRunTestWithLargeTimeout() = runTest(dispatchTimeoutMs = 5000) {
+    fun testRunTestWithLargeDispatchTimeout() = runTest(dispatchTimeoutMs = 5000) {
+        withContext(Dispatchers.Default) {
+            delay(50)
+        }
+    }
+
+    /** Tests that delays can be accounted for with a large enough timeout. */
+    @Test
+    fun testRunTestWithLargeTimeout() = runTest(timeout = 5000.milliseconds) {
         withContext(Dispatchers.Default) {
             delay(50)
         }
@@ -153,13 +209,13 @@ class RunTestTest {
         } catch (e: UncompletedCoroutinesError) {
             @Suppress("INVISIBLE_MEMBER")
             val suppressed = unwrap(e).suppressedExceptions
-            assertEquals(1, suppressed.size)
+            assertEquals(1, suppressed.size, "$suppressed")
             assertIs<TestException>(suppressed[0]).also {
                 assertEquals("A", it.message)
             }
         }
     }) {
-        runTest(dispatchTimeoutMs = 1) {
+        runTest(timeout = 1.milliseconds) {
             coroutineContext[CoroutineExceptionHandler]!!.handleException(coroutineContext, TestException("A"))
             withContext(Dispatchers.Default) {
                 delay(10000)
@@ -324,7 +380,7 @@ class RunTestTest {
         }
     }
 
-    /** Tests that [TestCoroutineScope.runTest] does not inherit the exception handler and works. */
+    /** Tests that [TestScope.runTest] does not inherit the exception handler and works. */
     @Test
     fun testScopeRunTestExceptionHandler(): TestResult {
         val scope = TestScope()
@@ -349,7 +405,7 @@ class RunTestTest {
      * The test will hang if this is not the case.
      */
     @Test
-    fun testCoroutineCompletingWithoutDispatch() = runTest(dispatchTimeoutMs = Long.MAX_VALUE) {
+    fun testCoroutineCompletingWithoutDispatch() = runTest(timeout = Duration.INFINITE) {
         launch(Dispatchers.Default) { delay(100) }
     }
 }

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

@@ -20,7 +20,7 @@ class StandardTestDispatcherTest: OrderedExecutionTestBase() {
     @AfterTest
     fun cleanup() {
         scope.runCurrent()
-        assertEquals(listOf(), scope.asSpecificImplementation().leave())
+        assertEquals(listOf(), scope.asSpecificImplementation().legacyLeave())
     }
 
     /** Tests that the [StandardTestDispatcher] follows an execution order similar to `runBlocking`. */