Add some tests, fix a bug

The bug was that the test would time out if every non-background
coroutine was waiting for something from the background ones.

Author: dkhalanskyjb

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

@@ -166,7 +166,7 @@ public fun TestScope.runTest(
     createTestResult {
         runTestCoroutine(it, dispatchTimeoutMs, TestScopeImpl::tryGetCompletionCause, testBody) {
             backgroundWorkScope.cancel()
-            testScheduler.advanceUntilIdle(backgroundIsIdle = false)
+            testScheduler.advanceUntilIdleOr { false }
             it.leave()
         }
     }
@@ -176,7 +176,7 @@ public fun TestScope.runTest(
  * Runs [testProcedure], creating a [TestResult].
  */
 @Suppress("NO_ACTUAL_FOR_EXPECT") // actually suppresses `TestResult`
-internal expect fun createTestResult(testProcedure: suspend () -> Unit): TestResult
+internal expect fun createTestResult(testProcedure: suspend CoroutineScope.() -> Unit): TestResult
 
 /** A coroutine context element indicating that the coroutine is running inside `runTest`. */
 internal object RunningInRunTest : CoroutineContext.Key<RunningInRunTest>, CoroutineContext.Element {
@@ -199,7 +199,7 @@ internal const val DEFAULT_DISPATCH_TIMEOUT_MS = 60_000L
  * 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>> runTestCoroutine(
+internal suspend fun <T: AbstractCoroutine<Unit>> CoroutineScope.runTestCoroutine(
     coroutine: T,
     dispatchTimeoutMs: Long,
     tryGetCompletionCause: T.() -> Throwable?,
@@ -220,16 +220,27 @@ internal suspend fun <T: AbstractCoroutine<Unit>> runTestCoroutine(
             completed = true
             continue
         }
-        select<Unit> {
-            coroutine.onJoin {
-                completed = true
+        // 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()
             }
-            scheduler.onDispatchEvent {
-                // we received knowledge that `scheduler` observed a dispatch event, so we reset the timeout
-            }
-            onTimeout(dispatchTimeoutMs) {
-                handleTimeout(coroutine, dispatchTimeoutMs, tryGetCompletionCause, cleanup)
+        }
+        try {
+            select<Unit> {
+                coroutine.onJoin {
+                    completed = true
+                }
+                scheduler.onDispatchEvent {
+                    // we received knowledge that `scheduler` observed a dispatch event, so we reset the timeout
+                }
+                onTimeout(dispatchTimeoutMs) {
+                    handleTimeout(coroutine, dispatchTimeoutMs, tryGetCompletionCause, cleanup)
+                }
             }
+        } finally {
+            backgroundWorkRunner.cancelAndJoin()
         }
     }
     coroutine.getCompletionExceptionOrNull()?.let { exception ->

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

@@ -96,7 +96,7 @@ private class UnconfinedTestDispatcherImpl(
     @Suppress("INVISIBLE_MEMBER")
     override fun dispatch(context: CoroutineContext, block: Runnable) {
         checkSchedulerInContext(scheduler, context)
-        scheduler.sendDispatchEvent()
+        scheduler.sendDispatchEvent(context)
 
         /** copy-pasted from [kotlinx.coroutines.Unconfined.dispatch] */
         /** It can only be called by the [yield] function. See also code of [yield] function. */

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

@@ -75,7 +75,7 @@ public class TestCoroutineScheduler : AbstractCoroutineContextElement(TestCorout
             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. */
-            sendDispatchEvent()
+            sendDispatchEvent(context)
             DisposableHandle {
                 synchronized(lock) {
                     events.remove(event)
@@ -86,11 +86,11 @@ public class TestCoroutineScheduler : AbstractCoroutineContextElement(TestCorout
 
     /**
      * Runs the next enqueued task, advancing the virtual time to the time of its scheduled awakening,
-     * unless all the remaining tasks are the background ones.
+     * unless [condition] holds.
      */
-    private fun tryRunNextTask(backgroundIsIdle: Boolean): Boolean {
+    internal fun tryRunNextTaskUnless(condition: () -> Boolean): Boolean {
         val event = synchronized(lock) {
-            if (backgroundIsIdle && events.none(TestDispatchEvent<*>::isForeground)) return false
+            if (condition()) return false
             val event = events.removeFirstOrNull() ?: return false
             if (currentTime > event.time)
                 currentTimeAheadOfEvents()
@@ -110,15 +110,14 @@ public class TestCoroutineScheduler : AbstractCoroutineContextElement(TestCorout
      * functionality, query [currentTime] before and after the execution to achieve the same result.
      */
     @ExperimentalCoroutinesApi
-    public fun advanceUntilIdle(): Unit = advanceUntilIdle(backgroundIsIdle = true)
+    public fun advanceUntilIdle(): Unit = advanceUntilIdleOr { events.none(TestDispatchEvent<*>::isForeground) }
 
     /**
-     * [backgroundIsIdle]: `true` if the background tasks should not be considered
-     * when checking if the scheduler is already idle.
+     * [condition]: guaranteed to be invoked under the lock.
      */
-    internal fun advanceUntilIdle(backgroundIsIdle: Boolean) {
+    internal fun advanceUntilIdleOr(condition: () -> Boolean) {
         while (true) {
-            if (!tryRunNextTask(backgroundIsIdle = backgroundIsIdle))
+            if (!tryRunNextTaskUnless(condition))
                 return
         }
     }
@@ -188,9 +187,12 @@ public class TestCoroutineScheduler : AbstractCoroutineContextElement(TestCorout
 
     /**
      * Notifies this scheduler about a dispatch event.
+     *
+     * [context] is the context in which the task will be dispatched.
      */
-    internal fun sendDispatchEvent() {
-        dispatchEvents.trySend(Unit)
+    internal fun sendDispatchEvent(context: CoroutineContext) {
+        if (context[BackgroundWork] !== BackgroundWork)
+            dispatchEvents.trySend(Unit)
     }
 
     /**
@@ -257,4 +259,4 @@ internal object BackgroundWork : CoroutineContext.Key<BackgroundWork>, Coroutine
 }
 
 private fun<T> ThreadSafeHeap<T>.none(predicate: (T) -> Boolean) where T: ThreadSafeHeapNode, T: Comparable<T> =
-    find(predicate) == null
+    find(predicate) == null

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

@@ -71,6 +71,7 @@ class RunTestTest {
 
     /** Tests that a dispatch timeout of `0` will fail the test if there are some dispatches outside the scheduler. */
     @Test
+    @NoNative // TODO: timeout leads to `Cannot execute task because event loop was shut down` on Native
     fun testRunTestWithZeroTimeoutWithUncontrolledDispatches() = testResultMap({ fn ->
         assertFailsWith<UncompletedCoroutinesError> { fn() }
     }) {

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

@@ -5,6 +5,8 @@
 package kotlinx.coroutines.test
 
 import kotlinx.coroutines.*
+import kotlinx.coroutines.channels.*
+import kotlinx.coroutines.flow.*
 import kotlin.coroutines.*
 import kotlin.test.*
 
@@ -327,6 +329,102 @@ class TestScopeTest {
         }
     }
 
+    /**
+     * Tests using [Flow.stateIn] as a background job.
+     */
+    @Test
+    fun testExampleBackgroundJob1() = runTest {
+        val myFlow = flow {
+            var i = 0
+            while (true) {
+                emit(++i)
+                delay(1)
+            }
+        }
+        val stateFlow = myFlow.stateIn(backgroundWorkScope, SharingStarted.Eagerly, 0)
+        var j = 0
+        repeat(100) {
+            assertEquals(j++, stateFlow.value)
+            delay(1)
+        }
+    }
+
+    /**
+     * A test from the documentation of [TestScope.backgroundWorkScope].
+     */
+    @Test
+    fun testExampleBackgroundJob2() = runTest {
+        val channel = Channel<Int>()
+        backgroundWorkScope.launch {
+            var i = 0
+            while (true) {
+                channel.send(i++)
+            }
+        }
+        repeat(100) {
+            assertEquals(it, channel.receive())
+        }
+    }
+
+    /**
+     * Tests that the test will timeout due to idleness even if some background tasks are running.
+     */
+    @Test
+    fun testBackgroundWorkNotPreventingTimeout(): TestResult = testResultMap({
+        try {
+            it()
+            fail("unreached")
+        } catch (_: UncompletedCoroutinesError) {
+
+        }
+    }) {
+        runTest(dispatchTimeoutMs = 100) {
+            backgroundWorkScope.launch {
+                while (true) {
+                    yield()
+                }
+            }
+            backgroundWorkScope.launch {
+                while (true) {
+                    delay(1)
+                }
+            }
+            val deferred = CompletableDeferred<Unit>()
+            deferred.await()
+        }
+
+    }
+
+    @Test
+    fun testUnconfinedBackgroundWorkNotPreventingTimeout(): TestResult = testResultMap({
+        try {
+            it()
+            fail("unreached")
+        } catch (_: UncompletedCoroutinesError) {
+
+        }
+    }) {
+        runTest(UnconfinedTestDispatcher(), dispatchTimeoutMs = 100) {
+            /**
+             * Having a coroutine like this will still cause the test to hang:
+                 backgroundWorkScope.launch {
+                     while (true) {
+                         yield()
+                     }
+                 }
+             * The reason is that even the initial [advanceUntilIdle] will never return in this case.
+             */
+            backgroundWorkScope.launch {
+                while (true) {
+                    delay(1)
+                }
+            }
+            val deferred = CompletableDeferred<Unit>()
+            deferred.await()
+        }
+
+    }
+
     companion object {
         internal val invalidContexts = listOf(
             Dispatchers.Default, // not a [TestDispatcher]

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

@@ -9,7 +9,7 @@ import kotlin.js.*
 @Suppress("ACTUAL_WITHOUT_EXPECT", "ACTUAL_TYPE_ALIAS_TO_CLASS_WITH_DECLARATION_SITE_VARIANCE")
 public actual typealias TestResult = Promise<Unit>
 
-internal actual fun createTestResult(testProcedure: suspend () -> Unit): TestResult =
+internal actual fun createTestResult(testProcedure: suspend CoroutineScope.() -> Unit): TestResult =
     GlobalScope.promise {
         testProcedure()
-    }
+    }

kotlinx-coroutines-test/jvm/src/TestBuildersJvm.kt

@@ -8,8 +8,8 @@ import kotlinx.coroutines.*
 @Suppress("ACTUAL_WITHOUT_EXPECT")
 public actual typealias TestResult = Unit
 
-internal actual fun createTestResult(testProcedure: suspend () -> Unit) {
+internal actual fun createTestResult(testProcedure: suspend CoroutineScope.() -> Unit) {
     runBlocking {
         testProcedure()
     }
-}
+}

kotlinx-coroutines-test/jvm/src/migration/TestBuildersDeprecated.kt

@@ -93,7 +93,7 @@ public fun runBlockingTestOnTestScope(
         null // the deferred was not completed yet; `scope.leave()` should complain then about unfinished jobs
     }
     scope.backgroundWorkScope.cancel()
-    scope.testScheduler.advanceUntilIdle(backgroundIsIdle = false)
+    scope.testScheduler.advanceUntilIdleOr { false }
     throwable?.let {
         val exceptions = try {
             scope.leave()

kotlinx-coroutines-test/jvm/src/migration/TestCoroutineDispatcher.kt

@@ -41,7 +41,7 @@ public class TestCoroutineDispatcher(public override val scheduler: TestCoroutin
     override fun dispatch(context: CoroutineContext, block: Runnable) {
         checkSchedulerInContext(scheduler, context)
         if (dispatchImmediately) {
-            scheduler.sendDispatchEvent()
+            scheduler.sendDispatchEvent(context)
             block.run()
         } else {
             post(block, context)

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

@@ -8,7 +8,7 @@ import kotlinx.coroutines.*
 @Suppress("ACTUAL_WITHOUT_EXPECT")
 public actual typealias TestResult = Unit
 
-internal actual fun createTestResult(testProcedure: suspend () -> Unit) {
+internal actual fun createTestResult(testProcedure: suspend CoroutineScope.() -> Unit) {
     runBlocking {
         testProcedure()
     }