Add TimeSource implementation to TestCoroutineScheduler

Author: hfhbd

kotlinx-coroutines-test/api/kotlinx-coroutines-test.api

@@ -65,6 +65,7 @@ public final class kotlinx/coroutines/test/TestCoroutineScheduler : kotlin/corou
 	public final fun advanceTimeBy (J)V
 	public final fun advanceUntilIdle ()V
 	public final fun getCurrentTime ()J
+	public final fun getTimeSource ()Lkotlin/time/TimeSource;
 	public final fun runCurrent ()V
 }
 
@@ -113,6 +114,7 @@ public final class kotlinx/coroutines/test/TestScopeKt {
 	public static final fun advanceTimeBy (Lkotlinx/coroutines/test/TestScope;J)V
 	public static final fun advanceUntilIdle (Lkotlinx/coroutines/test/TestScope;)V
 	public static final fun getCurrentTime (Lkotlinx/coroutines/test/TestScope;)J
+	public static final fun getTestTimeSource (Lkotlinx/coroutines/test/TestScope;)Lkotlin/time/TimeSource;
 	public static final fun runCurrent (Lkotlinx/coroutines/test/TestScope;)V
 }
 

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

@@ -12,6 +12,7 @@ import kotlinx.coroutines.internal.*
 import kotlinx.coroutines.selects.*
 import kotlin.coroutines.*
 import kotlin.jvm.*
+import kotlin.time.*
 
 /**
  * This is a scheduler for coroutines used in tests, providing the delay-skipping behavior.
@@ -26,7 +27,6 @@ import kotlin.jvm.*
  * haven't yet been dispatched (via [runCurrent]).
  */
 @ExperimentalCoroutinesApi
-// TODO: maybe make this a `TimeSource`?
 public class TestCoroutineScheduler : AbstractCoroutineContextElement(TestCoroutineScheduler),
     CoroutineContext.Element {
 
@@ -43,7 +43,7 @@ public class TestCoroutineScheduler : AbstractCoroutineContextElement(TestCorout
     /** This counter establishes some order on the events that happen at the same virtual time. */
     private val count = atomic(0L)
 
-    /** The current virtual time. */
+    /** The current virtual time in milliseconds. */
     @ExperimentalCoroutinesApi
     public var currentTime: Long = 0
         get() = synchronized(lock) { field }
@@ -193,6 +193,15 @@ public class TestCoroutineScheduler : AbstractCoroutineContextElement(TestCorout
      * Consumes the knowledge that a dispatch event happened recently.
      */
     internal val onDispatchEvent: SelectClause1<Unit> get() = dispatchEvents.onReceive
+
+    /**
+     * Returns the [TimeSource] representation of the virtual time of this scheduler.
+     */
+    @ExperimentalCoroutinesApi
+    @ExperimentalTime
+    public val timeSource: TimeSource = object : AbstractLongTimeSource(DurationUnit.MILLISECONDS) {
+        override fun read(): Long = currentTime
+    }
 }
 
 // Some error-throwing functions for pretty stack traces

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

@@ -7,6 +7,7 @@ package kotlinx.coroutines.test
 import kotlinx.coroutines.*
 import kotlinx.coroutines.internal.*
 import kotlin.coroutines.*
+import kotlin.time.*
 
 /**
  * A coroutine scope that for launching test coroutines.
@@ -84,6 +85,14 @@ public fun TestScope.runCurrent(): Unit = testScheduler.runCurrent()
 @ExperimentalCoroutinesApi
 public fun TestScope.advanceTimeBy(delayTimeMillis: Long): Unit = testScheduler.advanceTimeBy(delayTimeMillis)
 
+/**
+ * The [test scheduler][TestScope.testScheduler] as a [TimeSource].
+ * @see TestCoroutineScheduler.timeSource
+ */
+@ExperimentalCoroutinesApi
+@ExperimentalTime
+public val TestScope.testTimeSource: TimeSource get() = testScheduler.timeSource
+
 /**
  * Creates a [TestScope].
  *
@@ -237,4 +246,4 @@ internal class UncaughtExceptionsBeforeTest : IllegalStateException(
  * Thrown when a test has completed and there are tasks that are not completed or cancelled.
  */
 @ExperimentalCoroutinesApi
-internal class UncompletedCoroutinesError(message: String) : AssertionError(message)
+internal class UncompletedCoroutinesError(message: String) : AssertionError(message)

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

@@ -6,6 +6,8 @@ package kotlinx.coroutines.test
 
 import kotlinx.coroutines.*
 import kotlin.test.*
+import kotlin.time.*
+import kotlin.time.Duration.Companion.seconds
 
 class TestCoroutineSchedulerTest {
     /** Tests that `TestCoroutineScheduler` attempts to detect if there are several instances of it. */
@@ -308,6 +310,16 @@ class TestCoroutineSchedulerTest {
         }
     }
 
+    @Test
+    @ExperimentalTime
+    fun testAdvanceTimeSource() = runTest {
+        val expected = 1.seconds
+        val actual = testTimeSource.measureTime {
+            delay(expected)
+        }
+        assertEquals(expected, actual)
+    }
+
     private fun forTestDispatchers(block: (TestDispatcher) -> Unit): Unit =
         @Suppress("DEPRECATION")
         listOf(