More tests for the scheduler

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

@@ -70,7 +70,7 @@ class TestCoroutineSchedulerTest {
 
     /** Tests the basic functionality of [TestCoroutineScheduler.advanceTimeBy]. */
     @Test
-    fun testAdvanceTimeBy() {
+    fun testAdvanceTimeBy() = assertRunsFast {
         val scheduler = TestCoroutineScheduler()
         val scope = TestCoroutineScope(scheduler)
         var stage = 1
@@ -96,26 +96,89 @@ class TestCoroutineSchedulerTest {
         scope.cleanupTestCoroutines()
     }
 
+    /** Tests the basic functionality of [TestCoroutineScheduler.runCurrent]. */
+    @Test
+    fun testRunCurrent() = runBlockingTest {
+        var stage = 0
+        launch {
+            delay(1)
+            ++stage
+            delay(1)
+            stage += 10
+        }
+        launch {
+            delay(1)
+            ++stage
+            delay(1)
+            stage += 10
+        }
+        testScheduler.advanceTimeBy(1)
+        assertEquals(0, stage)
+        runCurrent()
+        assertEquals(2, stage)
+        testScheduler.advanceTimeBy(1)
+        assertEquals(2, stage)
+        runCurrent()
+        assertEquals(22, stage)
+    }
+
     /** Tests that [TestCoroutineScheduler.runCurrent] will not run new tasks after the current time has advanced. */
     @Test
-    fun testRunCurrentNotDrainingQueue() {
+    fun testRunCurrentNotDrainingQueue() = assertRunsFast {
         val scheduler = TestCoroutineScheduler()
         val scope = TestCoroutineScope(scheduler)
         var stage = 1
         scope.launch {
-            delay(1)
+            delay(SLOW)
             launch {
-                delay(1)
+                delay(SLOW)
                 stage = 3
             }
-            scheduler.advanceTimeBy(1)
+            scheduler.advanceTimeBy(SLOW)
             stage = 2
         }
-        scheduler.advanceTimeBy(1)
+        scheduler.advanceTimeBy(SLOW)
         assertEquals(1, stage)
         scheduler.runCurrent()
         assertEquals(2, stage)
         scheduler.runCurrent()
         assertEquals(3, stage)
     }
+
+    /** Tests that [TestCoroutineScheduler.advanceUntilIdle] doesn't hang when itself running in a scheduler task. */
+    @Test
+    fun testNestedAdvanceUntilIdle() = assertRunsFast {
+        val scheduler = TestCoroutineScheduler()
+        val scope = TestCoroutineScope(scheduler)
+        var executed = false
+        scope.launch {
+            launch {
+                delay(SLOW)
+                executed = true
+            }
+            scheduler.advanceUntilIdle()
+        }
+        scheduler.advanceUntilIdle()
+        assertTrue(executed)
+    }
+
+    /** Tests [yield] scheduling tasks for future execution and not executing immediately. */
+    @Test
+    fun testYield() {
+        val scope = TestCoroutineScope()
+        var stage = 0
+        scope.launch {
+            yield()
+            assertEquals(1, stage)
+            stage = 2
+        }
+        scope.launch {
+            yield()
+            assertEquals(2, stage)
+            stage = 3
+        }
+        assertEquals(0, stage)
+        stage = 1
+        scope.runCurrent()
+    }
 }

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

@@ -4,22 +4,27 @@
 
 package kotlinx.coroutines.test
 
-import kotlinx.coroutines.*
 import kotlin.test.*
 import kotlin.time.*
 
-const val SLOW = 10_000L
+/**
+ * The number of milliseconds that is sure not to pass [assertRunsFast].
+ */
+const val SLOW = 100_000L
 
 /**
- * Assert a block completes within a second or fail the suite
+ * Asserts that a block completed within [timeout].
  */
 @OptIn(ExperimentalTime::class)
-suspend fun CoroutineScope.assertRunsFast(block: suspend CoroutineScope.() -> Unit) {
-    val start = TimeSource.Monotonic.markNow()
-    // don't need to be fancy with timeouts here since anything longer than a few ms is an error
-    block()
-    val duration = start.elapsedNow()
-    assertTrue("All tests must complete within 2000ms (use longer timeouts to cause failure)") {
-        duration.inWholeSeconds < 2
-    }
+inline fun <T> assertRunsFast(timeout: Duration, block: () -> T): T {
+    val result: T
+    val elapsed = TimeSource.Monotonic.measureTime { result = block() }
+    assertTrue("Should complete in $timeout, but took $elapsed") { elapsed < timeout }
+    return result
 }
+
+/**
+ * Asserts that a block completed within two seconds.
+ */
+@OptIn(ExperimentalTime::class)
+inline fun <T> assertRunsFast(block: () -> T): T = assertRunsFast(Duration.seconds(2), block)