963: Provide optional collectionContext for WorkflowLayout.take

Closes #963

workflow-ui/core-android/src/main/java/com/squareup/workflow1/ui/WorkflowLayout.kt

@@ -17,13 +17,16 @@ import androidx.lifecycle.Lifecycle.State
 import androidx.lifecycle.Lifecycle.State.STARTED
 import androidx.lifecycle.coroutineScope
 import androidx.lifecycle.repeatOnLifecycle
+import kotlinx.coroutines.CoroutineDispatcher
 import kotlinx.coroutines.CoroutineScope
 import kotlinx.coroutines.Dispatchers
 import kotlinx.coroutines.Job
 import kotlinx.coroutines.flow.Flow
 import kotlinx.coroutines.flow.launchIn
 import kotlinx.coroutines.flow.onEach
 import kotlinx.coroutines.launch
+import kotlin.coroutines.CoroutineContext
+import kotlin.coroutines.EmptyCoroutineContext
 
 /**
  * A view that can be driven by a stream of [Screen] renderings passed to its [take] method.
@@ -86,15 +89,26 @@ public class WorkflowLayout(
    * Typically this comes from `ComponentActivity.lifecycle` or  `Fragment.lifecycle`.
    * @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.
    */
+  @OptIn(ExperimentalStdlibApi::class)
   public fun take(
     lifecycle: Lifecycle,
     renderings: Flow<Screen>,
-    repeatOnLifecycle: State = STARTED
+    repeatOnLifecycle: State = STARTED,
+    collectionContext: CoroutineContext = EmptyCoroutineContext
   ) {
+    // 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
-    lifecycle.coroutineScope.launch {
+    lifecycle.coroutineScope.launch(contextWithoutDispatcher) {
       lifecycle.repeatOnLifecycle(repeatOnLifecycle) {
+        require(coroutineContext[CoroutineDispatcher.Key] == lifecycleDispatcher) {
+          "Collection dispatch should happen on the lifecycle's dispatcher."
+        }
         renderings.collect { show(it) }
       }
     }

workflow-ui/core-android/src/test/java/com/squareup/workflow1/ui/WorkflowLayoutTest.kt

@@ -5,17 +5,24 @@ import android.os.Bundle
 import android.os.Parcelable
 import android.util.SparseArray
 import android.view.View
+import androidx.lifecycle.Lifecycle
+import androidx.lifecycle.testing.TestLifecycleOwner
 import androidx.test.core.app.ApplicationProvider
 import com.google.common.truth.Truth.assertThat
+import com.squareup.workflow1.ui.container.WrappedScreen
+import kotlinx.coroutines.ExperimentalCoroutinesApi
+import kotlinx.coroutines.flow.flowOf
+import kotlinx.coroutines.test.UnconfinedTestDispatcher
 import org.junit.Test
 import org.junit.runner.RunWith
 import org.robolectric.RobolectricTestRunner
 import org.robolectric.annotation.Config
+import kotlin.coroutines.CoroutineContext
 
 @RunWith(RobolectricTestRunner::class)
 // SDK 28 required for the four-arg constructor we use in our custom view classes.
 @Config(manifest = Config.NONE, sdk = [28])
-@OptIn(WorkflowUiExperimentalApi::class)
+@OptIn(WorkflowUiExperimentalApi::class, ExperimentalCoroutinesApi::class)
 internal class WorkflowLayoutTest {
   private val context: Context = ApplicationProvider.getApplicationContext()
 
@@ -38,4 +45,21 @@ internal class WorkflowLayoutTest {
     workflowLayout.restoreHierarchyState(viewState)
     // No crash, no bug.
   }
+
+  @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.
+  }
 }