Collect renderings on context specified in WorkflowLayout.take

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)
 }

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<TestScreen> {
+    override val viewFactory =
+      ScreenViewFactory.fromCode<TestScreen> { _, initialEnvironment, context, _ ->
+        ScreenViewHolder(initialEnvironment, View(context)) { _, _ -> }
+      }
+  }
+}

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) }
       }
     }
   }

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<Screen>()
+
+    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<TestContextElement>
+
+    @Volatile
+    var wasUsed = false
+  }
 }