From 15788a32fe22728e8742193090bbaad68bf4c521 Mon Sep 17 00:00:00 2001
From: Dmitry Khalanskiy <Dmitry.Khalanskiy@jetbrains.com>
Date: Fri, 15 Oct 2021 12:54:28 +0300
Subject: [PATCH] More tests for the scheduler

---
 .../common/test/TestCoroutineSchedulerTest.kt | 75 +++++++++++++++++--
 .../common/test/TestModuleHelpers.kt          |  9 +--
 2 files changed, 73 insertions(+), 11 deletions(-)

diff --git a/kotlinx-coroutines-test/common/test/TestCoroutineSchedulerTest.kt b/kotlinx-coroutines-test/common/test/TestCoroutineSchedulerTest.kt
index dfeabd530d..ab6ced4741 100644
--- a/kotlinx-coroutines-test/common/test/TestCoroutineSchedulerTest.kt
+++ b/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()
+    }
 }
diff --git a/kotlinx-coroutines-test/common/test/TestModuleHelpers.kt b/kotlinx-coroutines-test/common/test/TestModuleHelpers.kt
index a34dbfd6c7..ca8ffc85c8 100644
--- a/kotlinx-coroutines-test/common/test/TestModuleHelpers.kt
+++ b/kotlinx-coroutines-test/common/test/TestModuleHelpers.kt
@@ -4,22 +4,21 @@
 
 package kotlinx.coroutines.test
 
-import kotlinx.coroutines.*
 import kotlin.test.*
 import kotlin.time.*
 
-const val SLOW = 10_000L
+const val SLOW = 100_000L
 
 /**
  * Assert a block completes within a second or fail the suite
  */
 @OptIn(ExperimentalTime::class)
-suspend fun CoroutineScope.assertRunsFast(block: suspend CoroutineScope.() -> Unit) {
+inline fun <T> assertRunsFast(block: () -> T): T {
     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 result = block()
     val duration = start.elapsedNow()
     assertTrue("All tests must complete within 2000ms (use longer timeouts to cause failure)") {
         duration.inWholeSeconds < 2
     }
+    return result
 }