From 0d757f24baa58e4cc434a0c7742a5a3a1f4bf632 Mon Sep 17 00:00:00 2001 From: Stephen Edwards Date: Wed, 8 Oct 2025 13:14:16 -0400 Subject: [PATCH] Collect renderings on context specified in WorkflowLayout.take --- workflow-ui/core-android/build.gradle.kts | 2 + .../ui/WorkflowLayoutInstrumentedTest.kt | 65 +++++++++++++++++++ .../squareup/workflow1/ui/WorkflowLayout.kt | 20 +++--- .../workflow1/ui/WorkflowLayoutTest.kt | 57 +++++++++++----- 4 files changed, 116 insertions(+), 28 deletions(-) create mode 100644 workflow-ui/core-android/src/androidTest/java/com/squareup/workflow1/ui/WorkflowLayoutInstrumentedTest.kt diff --git a/workflow-ui/core-android/build.gradle.kts b/workflow-ui/core-android/build.gradle.kts index 1d489d7e81..7d15369a63 100644 --- a/workflow-ui/core-android/build.gradle.kts +++ b/workflow-ui/core-android/build.gradle.kts @@ -52,4 +52,6 @@ dependencies { testImplementation(libs.robolectric) testImplementation(libs.robolectric.annotations) testImplementation(libs.truth) + + androidTestImplementation(libs.androidx.lifecycle.testing) } diff --git a/workflow-ui/core-android/src/androidTest/java/com/squareup/workflow1/ui/WorkflowLayoutInstrumentedTest.kt b/workflow-ui/core-android/src/androidTest/java/com/squareup/workflow1/ui/WorkflowLayoutInstrumentedTest.kt new file mode 100644 index 0000000000..d24d2b09f6 --- /dev/null +++ b/workflow-ui/core-android/src/androidTest/java/com/squareup/workflow1/ui/WorkflowLayoutInstrumentedTest.kt @@ -0,0 +1,65 @@ +package com.squareup.workflow1.ui + +import android.view.View +import androidx.lifecycle.Lifecycle +import androidx.lifecycle.testing.TestLifecycleOwner +import androidx.test.platform.app.InstrumentationRegistry +import com.google.common.truth.Truth.assertThat +import kotlinx.coroutines.Dispatchers +import kotlinx.coroutines.ExperimentalCoroutinesApi +import kotlinx.coroutines.flow.MutableStateFlow +import org.junit.Test +import java.util.concurrent.CountDownLatch +import java.util.concurrent.TimeUnit + +/** + * Instrumented tests for [WorkflowLayout] that require a real Android environment. + * These tests verify behavior that cannot be properly tested with Robolectric, + * such as the main thread requirement for collecting renderings. + */ +@OptIn(ExperimentalCoroutinesApi::class) +internal class WorkflowLayoutInstrumentedTest { + + @Test + fun throwsWhenCollectingOnBackgroundThread() { + + var exception: Throwable? = null + val countDownLatch = CountDownLatch(1) + + Thread.setDefaultUncaughtExceptionHandler { _, throwable -> + exception = throwable + countDownLatch.countDown() + } + + val testLifeCycleOwner = TestLifecycleOwner(Lifecycle.State.CREATED) + val renderings = MutableStateFlow(TestScreen()) + + val nonMainThreadDispatcher = Dispatchers.IO + + val workflowLayout = WorkflowLayout(InstrumentationRegistry.getInstrumentation().context) + workflowLayout.take( + lifecycle = testLifeCycleOwner.lifecycle, + renderings = renderings, + collectionContext = nonMainThreadDispatcher + ) + + // start the lifecycle. + testLifeCycleOwner.lifecycle.currentState = Lifecycle.State.STARTED + + countDownLatch.await(1000, TimeUnit.MILLISECONDS) + + assertThat(exception).isNotNull() + assertThat(exception).isInstanceOf(IllegalArgumentException::class.java) + assertThat(exception?.message).contains("Collection dispatch must happen on the main thread!") + } + + /** + * Simple test screen for instrumented tests. + */ + private class TestScreen : AndroidScreen { + override val viewFactory = + ScreenViewFactory.fromCode { _, initialEnvironment, context, _ -> + ScreenViewHolder(initialEnvironment, View(context)) { _, _ -> } + } + } +} diff --git a/workflow-ui/core-android/src/main/java/com/squareup/workflow1/ui/WorkflowLayout.kt b/workflow-ui/core-android/src/main/java/com/squareup/workflow1/ui/WorkflowLayout.kt index 790796c680..d0f24f6aba 100644 --- a/workflow-ui/core-android/src/main/java/com/squareup/workflow1/ui/WorkflowLayout.kt +++ b/workflow-ui/core-android/src/main/java/com/squareup/workflow1/ui/WorkflowLayout.kt @@ -2,6 +2,7 @@ package com.squareup.workflow1.ui import android.content.Context import android.os.Build.VERSION +import android.os.Looper import android.os.Parcel import android.os.Parcelable import android.os.Parcelable.Creator @@ -16,10 +17,10 @@ import androidx.lifecycle.coroutineScope import androidx.lifecycle.repeatOnLifecycle import com.squareup.workflow1.ui.androidx.OnBackPressedDispatcherOwnerKey import com.squareup.workflow1.ui.androidx.WorkflowAndroidXSupport.onBackPressedDispatcherOwnerOrNull -import kotlinx.coroutines.CoroutineDispatcher import kotlinx.coroutines.Job import kotlinx.coroutines.flow.Flow import kotlinx.coroutines.launch +import kotlinx.coroutines.withContext import kotlin.coroutines.CoroutineContext import kotlin.coroutines.EmptyCoroutineContext @@ -89,8 +90,8 @@ public class WorkflowLayout( * @param [repeatOnLifecycle] the lifecycle state in which renderings should be actively * updated. Defaults to STARTED, which is appropriate for Activity and Fragment. * @param [collectionContext] additional [CoroutineContext] we want for the coroutine that is - * launched to collect the renderings. This should not override the [CoroutineDispatcher][kotlinx.coroutines.CoroutineDispatcher] - * but may include some other instrumentation elements. + * launched to collect the renderings, can include a different dispatcher - but we verify that + * it is a main thread dispatcher since we are updating views here! * * @return the [Job] started to collect [renderings], to give callers the option to * [cancel][Job.cancel] collection -- e.g., before calling [take] again with a new @@ -104,16 +105,15 @@ public class WorkflowLayout( repeatOnLifecycle: State = STARTED, collectionContext: CoroutineContext = EmptyCoroutineContext ): Job { - // We remove the dispatcher as we want to use what is provided by the lifecycle.coroutineScope. - val contextWithoutDispatcher = collectionContext.minusKey(CoroutineDispatcher.Key) - val lifecycleDispatcher = lifecycle.coroutineScope.coroutineContext[CoroutineDispatcher.Key] // Just like https://medium.com/androiddevelopers/a-safer-way-to-collect-flows-from-android-uis-23080b1f8bda - return lifecycle.coroutineScope.launch(contextWithoutDispatcher) { + return lifecycle.coroutineScope.launch { lifecycle.repeatOnLifecycle(repeatOnLifecycle) { - require(coroutineContext[CoroutineDispatcher.Key] == lifecycleDispatcher) { - "Collection dispatch should happen on the lifecycle's dispatcher." + withContext(collectionContext) { + require(Looper.myLooper() == Looper.getMainLooper()) { + "Collection dispatch must happen on the main thread!" + } + renderings.collect { show(it) } } - renderings.collect { show(it) } } } } diff --git a/workflow-ui/core-android/src/test/java/com/squareup/workflow1/ui/WorkflowLayoutTest.kt b/workflow-ui/core-android/src/test/java/com/squareup/workflow1/ui/WorkflowLayoutTest.kt index a0310243ad..0c919a49c7 100644 --- a/workflow-ui/core-android/src/test/java/com/squareup/workflow1/ui/WorkflowLayoutTest.kt +++ b/workflow-ui/core-android/src/test/java/com/squareup/workflow1/ui/WorkflowLayoutTest.kt @@ -17,14 +17,16 @@ import com.squareup.workflow1.ui.androidx.OnBackPressedDispatcherOwnerKey import com.squareup.workflow1.ui.navigation.WrappedScreen import kotlinx.coroutines.ExperimentalCoroutinesApi import kotlinx.coroutines.flow.MutableSharedFlow -import kotlinx.coroutines.flow.flowOf +import kotlinx.coroutines.flow.onEach import kotlinx.coroutines.test.UnconfinedTestDispatcher import kotlinx.coroutines.test.runTest import org.junit.Test import org.junit.runner.RunWith import org.robolectric.RobolectricTestRunner import org.robolectric.annotation.Config +import kotlin.coroutines.AbstractCoroutineContextElement import kotlin.coroutines.CoroutineContext +import kotlin.coroutines.coroutineContext @RunWith(RobolectricTestRunner::class) // SDK 28 required for the four-arg constructor we use in our custom view classes. @@ -92,23 +94,6 @@ internal class WorkflowLayoutTest { unoriginal.show(BScreen(), env) } - @Test fun usesLifecycleDispatcher() { - val lifecycleDispatcher = UnconfinedTestDispatcher() - val collectionContext: CoroutineContext = UnconfinedTestDispatcher() - val testLifecycle = TestLifecycleOwner( - Lifecycle.State.RESUMED, - lifecycleDispatcher - ) - - workflowLayout.take( - lifecycle = testLifecycle.lifecycle, - renderings = flowOf(WrappedScreen(), WrappedScreen()), - collectionContext = collectionContext - ) - - // No crash then we safely removed the dispatcher. - } - @Test fun takes() { val lifecycleDispatcher = UnconfinedTestDispatcher() val testLifecycle = TestLifecycleOwner( @@ -147,6 +132,35 @@ internal class WorkflowLayoutTest { } } + @Test fun usesProvidedCoroutineContext() { + val lifecycleDispatcher = UnconfinedTestDispatcher() + val testLifecycle = TestLifecycleOwner( + initialState = Lifecycle.State.RESUMED, + coroutineDispatcher = lifecycleDispatcher + ) + val flow = MutableSharedFlow() + + val testElement = TestContextElement() + + val trackedFlow = flow.onEach { + if (coroutineContext[TestContextElement] != null) { + testElement.wasUsed = true + } + } + + runTest(lifecycleDispatcher) { + workflowLayout.take( + lifecycle = testLifecycle.lifecycle, + renderings = trackedFlow, + collectionContext = testElement + ) + + flow.emit(WrappedScreen()) + + assertThat(testElement.wasUsed).isTrue() + } + } + private class BundleSavingView(context: Context) : View(context) { var saved = false @@ -166,4 +180,11 @@ internal class WorkflowLayoutTest { id = 42 } } + + private class TestContextElement : AbstractCoroutineContextElement(Key) { + companion object Key : CoroutineContext.Key + + @Volatile + var wasUsed = false + } }