diff --git a/kotlin/workflow-runtime/src/main/java/com/squareup/workflow/internal/RealWorkflowContext.kt b/kotlin/workflow-runtime/src/main/java/com/squareup/workflow/internal/RealWorkflowContext.kt index 70fa31165..1a72e75fa 100644 --- a/kotlin/workflow-runtime/src/main/java/com/squareup/workflow/internal/RealWorkflowContext.kt +++ b/kotlin/workflow-runtime/src/main/java/com/squareup/workflow/internal/RealWorkflowContext.kt @@ -48,8 +48,12 @@ internal class RealWorkflowContext( private val childCases = mutableListOf>() private val teardownHooks = mutableListOf<() -> Unit>() + /** Used to prevent modifications to this object after [buildBehavior] is called. */ + private var frozen = false + override fun onEvent(handler: (EventT) -> WorkflowAction): EventHandler { + 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( key: String, handler: (ChannelUpdate) -> WorkflowAction ) { + checkNotFrozen() subscriptionCases += SubscriptionCase(channelProvider, Pair(type, key), handler) } @@ -82,6 +87,7 @@ internal class RealWorkflowContext( handler: (ChildOutputT) -> WorkflowAction ): ChildRenderingT { // @formatter:on + checkNotFrozen() val id = child.id(key) val case: WorkflowOutputCase = WorkflowOutputCase(child, id, input, handler) @@ -90,16 +96,25 @@ internal class RealWorkflowContext( } override fun onTeardown(handler: () -> Unit) { + checkNotFrozen() teardownHooks += handler } /** * Constructs an immutable [Behavior] from the context. */ - fun buildBehavior(): Behavior = Behavior( - childCases = childCases.toList(), - subscriptionCases = subscriptionCases.toList(), - nextActionFromEvent = nextUpdateFromEvent, - teardownHooks = teardownHooks.toList() - ) + fun buildBehavior(): Behavior { + 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." + } } diff --git a/kotlin/workflow-runtime/src/test/java/com/squareup/workflow/internal/RealWorkflowContextTest.kt b/kotlin/workflow-runtime/src/test/java/com/squareup/workflow/internal/RealWorkflowContextTest.kt index 2709b7dae..144882f35 100644 --- a/kotlin/workflow-runtime/src/test/java/com/squareup/workflow/internal/RealWorkflowContextTest.kt +++ b/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(PoisonComposer()) val expectedUpdate = noop() val handler = context.onEvent { 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 { context.onEvent { fail() } } + assertFailsWith { context.onTeardown { fail() } } + assertFailsWith { + context.onReceive({ fail() }, Unit::class.starProjectedType) { fail() } + } + val child = Workflow.stateless { fail() } + assertFailsWith { context.compose(child) } + assertFailsWith { context.buildBehavior() } + } }