Fix 'TestCoroutineScheduler' being ignored on 'TestCoroutineScope' creation

Author: dkhalanskyjb

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

@@ -96,7 +96,7 @@ internal fun CoroutineContext.checkTestScopeArguments(): Pair<CoroutineContext,
             dispatcher
         }
         null -> {
-            scheduler = TestCoroutineScheduler()
+            scheduler = get(TestCoroutineScheduler) ?: TestCoroutineScheduler()
             TestCoroutineDispatcher(scheduler)
         }
         else -> throw IllegalArgumentException("Dispatcher must implement TestDispatcher: $dispatcher")

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

@@ -5,21 +5,54 @@
 package kotlinx.coroutines.test
 
 import kotlinx.coroutines.*
+import kotlin.coroutines.*
 import kotlin.test.*
 
 class TestCoroutineScopeTest {
+    /** Tests failing to create a [TestCoroutineScope] with incorrect contexts. */
     @Test
-    fun whenGivenInvalidExceptionHandler_throwsException() {
-        val handler = CoroutineExceptionHandler {  _, _ -> }
-        assertFails {
-            TestCoroutineScope(handler)
+    fun testCreateThrowsOnInvalidArguments() {
+        for (ctx in invalidContexts) {
+            assertFailsWith<IllegalArgumentException> {
+                TestCoroutineScope(ctx)
+            }
         }
     }
 
+    /** Tests that a newly-created [TestCoroutineScope] provides the correct scheduler. */
     @Test
-    fun whenGivenInvalidDispatcher_throwsException() {
-        assertFails {
-            TestCoroutineScope(Dispatchers.Default)
+    fun testCreateProvidesScheduler() {
+        // Creates a new scheduler.
+        run {
+            val scope = TestCoroutineScope()
+            assertNotNull(scope.coroutineContext[TestCoroutineScheduler])
+        }
+        // Reuses the scheduler that the dispatcher is linked to.
+        run {
+            val dispatcher = TestCoroutineDispatcher()
+            val scope = TestCoroutineScope(dispatcher)
+            assertSame(dispatcher.scheduler, scope.coroutineContext[TestCoroutineScheduler])
+        }
+        // Uses the scheduler passed to it.
+        run {
+            val scheduler = TestCoroutineScheduler()
+            val scope = TestCoroutineScope(scheduler)
+            assertSame(scheduler, scope.coroutineContext[TestCoroutineScheduler])
+            assertSame(scheduler, (scope.coroutineContext[ContinuationInterceptor] as TestDispatcher).scheduler)
+        }
+        // Doesn't touch the passed dispatcher and the scheduler if they match.
+        run {
+            val scheduler = TestCoroutineScheduler()
+            val dispatcher = TestCoroutineDispatcher(scheduler)
+            val scope = TestCoroutineScope(scheduler + dispatcher)
+            assertSame(scheduler, scope.coroutineContext[TestCoroutineScheduler])
+            assertSame(dispatcher, scope.coroutineContext[ContinuationInterceptor])
         }
     }
+
+    private val invalidContexts = listOf(
+        Dispatchers.Default, // not a [TestDispatcher]
+        TestCoroutineDispatcher() + TestCoroutineScheduler(), // the dispatcher is not linked to the scheduler
+        CoroutineExceptionHandler { _, _ -> }, // not an `UncaughtExceptionCaptor`
+    )
 }