Throw if WorkflowContext is accessed after compose returns.

kotlin/workflow-runtime/src/main/java/com/squareup/workflow/internal/RealWorkflowContext.kt

@@ -48,8 +48,12 @@ internal class RealWorkflowContext<StateT : Any, OutputT : Any>(
   private val childCases = mutableListOf<WorkflowOutputCase<*, *, StateT, OutputT>>()
   private val teardownHooks = mutableListOf<() -> Unit>()
 
+  /** Used to prevent modifications to this object after [buildBehavior] is called. */
+  private var frozen = false
+
   override fun <EventT : Any> onEvent(handler: (EventT) -> WorkflowAction<StateT, OutputT>):
       EventHandler<EventT> {
+    checkNotFrozen()
     return EventHandler { event ->
       // Run the handler synchronously, so we only have to emit the resulting action and don't need the
       // update channel to be generic on each event type.
@@ -70,6 +74,7 @@ internal class RealWorkflowContext<StateT : Any, OutputT : Any>(
     key: String,
     handler: (ChannelUpdate<E>) -> WorkflowAction<StateT, OutputT>
   ) {
+    checkNotFrozen()
     subscriptionCases += SubscriptionCase(channelProvider, Pair(type, key), handler)
   }
 
@@ -82,6 +87,7 @@ internal class RealWorkflowContext<StateT : Any, OutputT : Any>(
         handler: (ChildOutputT) -> WorkflowAction<StateT, OutputT>
       ): ChildRenderingT {
     // @formatter:on
+    checkNotFrozen()
     val id = child.id(key)
     val case: WorkflowOutputCase<ChildInputT, ChildOutputT, StateT, OutputT> =
       WorkflowOutputCase(child, id, input, handler)
@@ -90,16 +96,25 @@ internal class RealWorkflowContext<StateT : Any, OutputT : Any>(
   }
 
   override fun onTeardown(handler: () -> Unit) {
+    checkNotFrozen()
     teardownHooks += handler
   }
 
   /**
    * Constructs an immutable [Behavior] from the context.
    */
-  fun buildBehavior(): Behavior<StateT, OutputT> = Behavior(
-      childCases = childCases.toList(),
-      subscriptionCases = subscriptionCases.toList(),
-      nextActionFromEvent = nextUpdateFromEvent,
-      teardownHooks = teardownHooks.toList()
-  )
+  fun buildBehavior(): Behavior<StateT, OutputT> {
+    checkNotFrozen()
+    frozen = true
+    return Behavior(
+        childCases = childCases.toList(),
+        subscriptionCases = subscriptionCases.toList(),
+        nextActionFromEvent = nextUpdateFromEvent,
+        teardownHooks = teardownHooks.toList()
+    )
+  }
+
+  private fun checkNotFrozen() = check(!frozen) {
+    "WorkflowContext cannot be used after compose method returns."
+  }
 }

kotlin/workflow-runtime/src/test/java/com/squareup/workflow/internal/RealWorkflowContextTest.kt

@@ -23,11 +23,15 @@ import com.squareup.workflow.Workflow
 import com.squareup.workflow.WorkflowAction.Companion.emitOutput
 import com.squareup.workflow.WorkflowAction.Companion.noop
 import com.squareup.workflow.WorkflowContext
+import com.squareup.workflow.compose
 import com.squareup.workflow.internal.Behavior.WorkflowOutputCase
 import com.squareup.workflow.internal.RealWorkflowContext.Composer
 import com.squareup.workflow.internal.RealWorkflowContextTest.TestComposer.Rendering
+import com.squareup.workflow.stateless
+import kotlin.reflect.full.starProjectedType
 import kotlin.test.Test
 import kotlin.test.assertEquals
+import kotlin.test.assertFailsWith
 import kotlin.test.assertFalse
 import kotlin.test.assertSame
 import kotlin.test.assertTrue
@@ -85,11 +89,11 @@ class RealWorkflowContextTest {
     val context = RealWorkflowContext<String, String>(PoisonComposer())
     val expectedUpdate = noop<String, String>()
     val handler = context.onEvent<String> { expectedUpdate }
-    assertFalse(context.buildBehavior().nextActionFromEvent.isCompleted)
+    val behavior = context.buildBehavior()
+    assertFalse(behavior.nextActionFromEvent.isCompleted)
 
     handler("")
 
-    val behavior = context.buildBehavior()
     assertTrue(behavior.nextActionFromEvent.isCompleted)
     val actualUpdate = behavior.nextActionFromEvent.getCompleted()
     assertSame(expectedUpdate, actualUpdate)
@@ -134,4 +138,18 @@ class RealWorkflowContextTest {
     assertEquals(1, childCases.size)
     assertSame(case, childCases.single())
   }
+
+  @Test fun `all methods throw after buildBehavior`() {
+    val context = RealWorkflowContext(TestComposer())
+    context.buildBehavior()
+
+    assertFailsWith<IllegalStateException> { context.onEvent<Unit> { fail() } }
+    assertFailsWith<IllegalStateException> { context.onTeardown { fail() } }
+    assertFailsWith<IllegalStateException> {
+      context.onReceive<Unit>({ fail() }, Unit::class.starProjectedType) { fail() }
+    }
+    val child = Workflow.stateless<Nothing, Unit> { fail() }
+    assertFailsWith<IllegalStateException> { context.compose(child) }
+    assertFailsWith<IllegalStateException> { context.buildBehavior() }
+  }
 }