Use the TestCoroutineDispatcher from Dispatchers.Main by default

Author: dkhalanskyjb

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

@@ -120,10 +120,12 @@ public expect class TestResult
  *
  * ### Delay-skipping
  *
- * Delay-skipping is achieved by using virtual time. [TestCoroutineScheduler] is automatically created (if it wasn't
- * passed in some way in [context]) and can be used to control the virtual time, advancing it, running the tasks
- * scheduled at a specific time etc. Some convenience methods are available on [TestCoroutineScope] to control the
- * scheduler.
+ * Delay-skipping is achieved by using virtual time.
+ * If [Dispatchers.Main] is set to a [TestDispatcher] via [Dispatchers.setMain] before the test,
+ * then its [TestCoroutineScheduler] is used;
+ * otherwise, a new one is automatically created (or taken from [context] in some way) and can be used to control
+ * the virtual time, advancing it, running the tasks scheduled at a specific time etc.
+ * Some convenience methods are available on [TestCoroutineScope] to control the scheduler.
  *
  * Delays in code that runs inside dispatchers that don't use a [TestCoroutineScheduler] don't get skipped:
  * ```

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

@@ -6,6 +6,8 @@ package kotlinx.coroutines.test
 
 import kotlinx.coroutines.*
 import kotlinx.coroutines.internal.*
+import kotlinx.coroutines.test.internal.*
+import kotlinx.coroutines.test.internal.TestMainDispatcher
 import kotlin.coroutines.*
 
 /**
@@ -140,7 +142,10 @@ public fun TestCoroutineScope(context: CoroutineContext = EmptyCoroutineContext)
  *
  * It ensures that all the test module machinery is properly initialized.
  * * If [context] doesn't define a [TestCoroutineScheduler] for orchestrating the virtual time used for delay-skipping,
- *   a new one is created, unless a [TestDispatcher] is provided, in which case [TestDispatcher.scheduler] is used.
+ *   a new one is created, unless either
+ *   - a [TestDispatcher] is provided, in which case [TestDispatcher.scheduler] is used;
+ *   - at the moment of the creation of the scope, [Dispatchers.Main] is delegated to a [TestDispatcher], in which case
+ *     its [TestCoroutineScheduler] is used.
  * * If [context] doesn't have a [ContinuationInterceptor], a [StandardTestDispatcher] is created.
  * * A [CoroutineExceptionHandler] is created that makes [TestCoroutineScope.cleanupTestCoroutines] throw if there were
  *   any uncaught exceptions, or forwards the exceptions further in a platform-specific manner if the cleanup was
@@ -173,7 +178,9 @@ public fun createTestCoroutineScope(context: CoroutineContext = EmptyCoroutineCo
             dispatcher
         }
         null -> {
-            scheduler = context[TestCoroutineScheduler] ?: TestCoroutineScheduler()
+            val mainDispatcherScheduler =
+                ((Dispatchers.Main as? TestMainDispatcher)?.delegate as? TestDispatcher)?.scheduler
+            scheduler = context[TestCoroutineScheduler] ?: mainDispatcherScheduler ?: TestCoroutineScheduler()
             StandardTestDispatcher(scheduler)
         }
         else -> throw IllegalArgumentException("Dispatcher must implement TestDispatcher: $dispatcher")

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

@@ -18,7 +18,7 @@ import kotlin.jvm.*
 @ExperimentalCoroutinesApi
 public fun Dispatchers.setMain(dispatcher: CoroutineDispatcher) {
     require(dispatcher !is TestMainDispatcher) { "Dispatchers.setMain(Dispatchers.Main) is prohibited, probably Dispatchers.resetMain() should be used instead" }
-    getTestMainDispatcher().setDispatcher(dispatcher)
+    getTestMainDispatcher().delegate = dispatcher
 }
 
 /**

kotlinx-coroutines-test/common/src/internal/TestMainDispatcher.kt

@@ -10,7 +10,7 @@ import kotlin.coroutines.*
  * The testable main dispatcher used by kotlinx-coroutines-test.
  * It is a [MainCoroutineDispatcher] that delegates all actions to a settable delegate.
  */
-internal class TestMainDispatcher(private var delegate: CoroutineDispatcher):
+internal class TestMainDispatcher(var delegate: CoroutineDispatcher):
     MainCoroutineDispatcher(),
     Delay by (delegate as? Delay ?: defaultDelay)
 {
@@ -25,10 +25,6 @@ internal class TestMainDispatcher(private var delegate: CoroutineDispatcher):
 
     override fun dispatchYield(context: CoroutineContext, block: Runnable) = delegate.dispatchYield(context, block)
 
-    fun setDispatcher(dispatcher: CoroutineDispatcher) {
-        delegate = dispatcher
-    }
-
     fun resetDispatcher() {
         delegate = mainDispatcher
     }

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

@@ -48,6 +48,25 @@ class TestCoroutineScopeTest {
             assertSame(scheduler, scope.coroutineContext[TestCoroutineScheduler])
             assertSame(dispatcher, scope.coroutineContext[ContinuationInterceptor])
         }
+        // Reuses the scheduler of `Dispatchers.Main`
+        run {
+            val scheduler = TestCoroutineScheduler()
+            val mainDispatcher = StandardTestDispatcher(scheduler)
+            Dispatchers.setMain(mainDispatcher)
+            val scope = createTestCoroutineScope()
+            assertSame(scheduler, scope.coroutineContext[TestCoroutineScheduler])
+            assertNotSame(mainDispatcher, scope.coroutineContext[ContinuationInterceptor])
+        }
+        // Does not reuse the scheduler of `Dispatchers.Main` if one is explicitly passed
+        run {
+            val mainDispatcher = StandardTestDispatcher()
+            Dispatchers.setMain(mainDispatcher)
+            val scheduler = TestCoroutineScheduler()
+            val scope = createTestCoroutineScope(scheduler)
+            assertSame(scheduler, scope.coroutineContext[TestCoroutineScheduler])
+            assertNotSame(mainDispatcher.scheduler, scope.coroutineContext[TestCoroutineScheduler])
+            assertNotSame(mainDispatcher, scope.coroutineContext[ContinuationInterceptor])
+        }
     }
 
     /** Tests that the cleanup procedure throws if there were uncompleted delays by the end. */