Return the Job from WorkflowLayout.take so it can be canceled.

workflow-ui/core-android/api/core-android.api

@@ -128,8 +128,8 @@ public final class com/squareup/workflow1/ui/WorkflowLayout : android/widget/Fra
 	public synthetic fun <init> (Landroid/content/Context;Landroid/util/AttributeSet;ILkotlin/jvm/internal/DefaultConstructorMarker;)V
 	public final fun show (Lcom/squareup/workflow1/ui/Screen;Lcom/squareup/workflow1/ui/ViewEnvironment;)V
 	public static synthetic fun show$default (Lcom/squareup/workflow1/ui/WorkflowLayout;Lcom/squareup/workflow1/ui/Screen;Lcom/squareup/workflow1/ui/ViewEnvironment;ILjava/lang/Object;)V
-	public final fun take (Landroidx/lifecycle/Lifecycle;Lkotlinx/coroutines/flow/Flow;Landroidx/lifecycle/Lifecycle$State;Lkotlin/coroutines/CoroutineContext;)V
-	public static synthetic fun take$default (Lcom/squareup/workflow1/ui/WorkflowLayout;Landroidx/lifecycle/Lifecycle;Lkotlinx/coroutines/flow/Flow;Landroidx/lifecycle/Lifecycle$State;Lkotlin/coroutines/CoroutineContext;ILjava/lang/Object;)V
+	public final fun take (Landroidx/lifecycle/Lifecycle;Lkotlinx/coroutines/flow/Flow;Landroidx/lifecycle/Lifecycle$State;Lkotlin/coroutines/CoroutineContext;)Lkotlinx/coroutines/Job;
+	public static synthetic fun take$default (Lcom/squareup/workflow1/ui/WorkflowLayout;Landroidx/lifecycle/Lifecycle;Lkotlinx/coroutines/flow/Flow;Landroidx/lifecycle/Lifecycle$State;Lkotlin/coroutines/CoroutineContext;ILjava/lang/Object;)Lkotlinx/coroutines/Job;
 }
 
 public final class com/squareup/workflow1/ui/WorkflowViewStub : android/view/View {

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

@@ -17,6 +17,7 @@ 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 kotlin.coroutines.CoroutineContext
@@ -91,19 +92,24 @@ public class WorkflowLayout(
    * @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.
+   *
+   * @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
+   * [renderings] flow. In most cases the caller can ignore this, interacting with
+   * the [Job] is very unusual.
    */
   @OptIn(ExperimentalStdlibApi::class)
   public fun take(
     lifecycle: Lifecycle,
     renderings: Flow<Screen>,
     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
-    lifecycle.coroutineScope.launch(contextWithoutDispatcher) {
+    return lifecycle.coroutineScope.launch(contextWithoutDispatcher) {
       lifecycle.repeatOnLifecycle(repeatOnLifecycle) {
         require(coroutineContext[CoroutineDispatcher.Key] == lifecycleDispatcher) {
           "Collection dispatch should happen on the lifecycle's dispatcher."

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

@@ -5,16 +5,21 @@ import android.os.Bundle
 import android.os.Parcelable
 import android.util.SparseArray
 import android.view.View
+import androidx.activity.OnBackPressedDispatcher
 import androidx.activity.OnBackPressedDispatcherOwner
+import androidx.activity.setViewTreeOnBackPressedDispatcherOwner
+import androidx.core.view.get
 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.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.test.UnconfinedTestDispatcher
+import kotlinx.coroutines.test.runTest
 import org.junit.Test
 import org.junit.runner.RunWith
 import org.robolectric.RobolectricTestRunner
@@ -28,7 +33,15 @@ import kotlin.coroutines.CoroutineContext
 internal class WorkflowLayoutTest {
   private val context: Context = ApplicationProvider.getApplicationContext()
 
-  private val workflowLayout = WorkflowLayout(context).apply { id = 42 }
+  private val workflowLayout = WorkflowLayout(context).apply {
+    id = 42
+    setViewTreeOnBackPressedDispatcherOwner(object : OnBackPressedDispatcherOwner {
+      override fun getOnBackPressedDispatcher(): OnBackPressedDispatcher {
+        error("yeah no")
+      }
+      override val lifecycle: Lifecycle get() = error("nope")
+    })
+  }
 
   @Test fun ignoresAlienViewState() {
     val weirdView = BundleSavingView(context)
@@ -91,6 +104,44 @@ internal class WorkflowLayoutTest {
     // No crash then we safely removed the dispatcher.
   }
 
+  @Test fun takes() {
+    val lifecycleDispatcher = UnconfinedTestDispatcher()
+    val testLifecycle = TestLifecycleOwner(
+      initialState = Lifecycle.State.RESUMED,
+      coroutineDispatcher = lifecycleDispatcher
+    )
+    val flow = MutableSharedFlow<Screen>()
+
+    runTest(lifecycleDispatcher) {
+      workflowLayout.take(
+        lifecycle = testLifecycle.lifecycle,
+        renderings = flow,
+      )
+      assertThat(workflowLayout[0]).isInstanceOf(WorkflowViewStub::class.java)
+      flow.emit(WrappedScreen())
+      assertThat(workflowLayout[0]).isNotInstanceOf(WorkflowViewStub::class.java)
+    }
+  }
+
+  @Test fun canStopTaking() {
+    val lifecycleDispatcher = UnconfinedTestDispatcher()
+    val testLifecycle = TestLifecycleOwner(
+      initialState = Lifecycle.State.RESUMED,
+      coroutineDispatcher = lifecycleDispatcher
+    )
+    val flow = MutableSharedFlow<Screen>()
+
+    runTest(lifecycleDispatcher) {
+      val job = workflowLayout.take(
+        lifecycle = testLifecycle.lifecycle,
+        renderings = flow,
+      )
+      job.cancel()
+      flow.emit(WrappedScreen())
+      assertThat(workflowLayout[0]).isInstanceOf(WorkflowViewStub::class.java)
+    }
+  }
+
   private class BundleSavingView(context: Context) : View(context) {
     var saved = false