2 files changed, 48 insertions, 3 deletions
diff --git a/packages/SystemUI/tests/src/com/android/systemui/coroutines/FlowTest.kt b/packages/SystemUI/tests/src/com/android/systemui/coroutines/FlowTest.kt
new file mode 100644
index 000000000000..1e4753e5b4a5
--- /dev/null
+++ b/packages/SystemUI/tests/src/com/android/systemui/coroutines/FlowTest.kt
@@ -0,0 +1,24 @@
+package com.android.systemui.coroutines
+
+import android.testing.AndroidTestingRunner
+import androidx.test.filters.SmallTest
+import com.android.systemui.SysuiTestCase
+import com.google.common.truth.Truth.assertThat
+import kotlinx.coroutines.ExperimentalCoroutinesApi
+import kotlinx.coroutines.flow.flowOf
+import kotlinx.coroutines.test.runTest
+import org.junit.Test
+import org.junit.runner.RunWith
+
+@OptIn(ExperimentalCoroutinesApi::class)
+@SmallTest
+@RunWith(AndroidTestingRunner::class)
+class FlowTest : SysuiTestCase() {
+
+    @Test
+    fun collectLastValue() = runTest {
+        val flow = flowOf(0, 1, 2)
+        val lastValue by collectLastValue(flow)
+        assertThat(lastValue).isEqualTo(2)
+    }
+}
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/coroutines/Flow.kt b/packages/SystemUI/tests/utils/src/com/android/systemui/coroutines/Flow.kt
index b7a8d2e9f684..9b4f4969f9e9 100644
--- a/packages/SystemUI/tests/utils/src/com/android/systemui/coroutines/Flow.kt
+++ b/packages/SystemUI/tests/utils/src/com/android/systemui/coroutines/Flow.kt
@@ -18,6 +18,8 @@ package com.android.systemui.coroutines
 
 import kotlin.coroutines.CoroutineContext
 import kotlin.coroutines.EmptyCoroutineContext
+import kotlin.properties.ReadOnlyProperty
+import kotlin.reflect.KProperty
 import kotlinx.coroutines.CoroutineStart
 import kotlinx.coroutines.Job
 import kotlinx.coroutines.flow.Flow
@@ -25,16 +27,35 @@ import kotlinx.coroutines.launch
 import kotlinx.coroutines.test.TestScope
 import kotlinx.coroutines.test.runCurrent
 
-/** Collect [flow] in a new [Job] and return a getter for the last collected value. */
+/**
+ * Collect [flow] in a new [Job] and return a getter for the last collected value.
+ * ```
+ * fun myTest() = runTest {
+ *   // ...
+ *   val actual by collectLastValue(underTest.flow)
+ *   assertThat(actual).isEqualTo(expected)
+ * }
+ * ```
+ */
 fun <T> TestScope.collectLastValue(
     flow: Flow<T>,
     context: CoroutineContext = EmptyCoroutineContext,
     start: CoroutineStart = CoroutineStart.DEFAULT,
-): () -> T? {
+): FlowValue<T?> {
     var lastValue: T? = null
     backgroundScope.launch(context, start) { flow.collect { lastValue = it } }
-    return {
+    return FlowValueImpl {
         runCurrent()
         lastValue
     }
 }
+
+/** @see collectLastValue */
+interface FlowValue<T> : ReadOnlyProperty<Any?, T?> {
+    operator fun invoke(): T?
+}
+
+private class FlowValueImpl<T>(private val block: () -> T?) : FlowValue<T> {
+    override operator fun invoke(): T? = block()
+    override fun getValue(thisRef: Any?, property: KProperty<*>): T? = invoke()
+}