Add testRender with CoroutineScope for SessionWorkflow

Fixes #1138

Author: steve-the-edwards

workflow-testing/src/main/java/com/squareup/workflow1/testing/RenderTester.kt

@@ -1,14 +1,17 @@
 package com.squareup.workflow1.testing
 
 import com.squareup.workflow1.ActionApplied
+import com.squareup.workflow1.SessionWorkflow
 import com.squareup.workflow1.StatefulWorkflow
 import com.squareup.workflow1.Workflow
 import com.squareup.workflow1.WorkflowAction
+import com.squareup.workflow1.WorkflowExperimentalApi
 import com.squareup.workflow1.WorkflowIdentifier
 import com.squareup.workflow1.WorkflowOutput
 import com.squareup.workflow1.identifier
 import com.squareup.workflow1.testing.RenderTester.ChildWorkflowMatch
 import com.squareup.workflow1.workflowIdentifier
+import kotlinx.coroutines.CoroutineScope
 import kotlin.reflect.KClass
 import kotlin.reflect.KType
 import kotlin.reflect.KTypeProjection
@@ -19,14 +22,43 @@ import kotlin.reflect.KTypeProjection
  *
  * See [RenderTester] for usage documentation.
  */
+@OptIn(WorkflowExperimentalApi::class) // Opt-in is only for the argument check.
 @Suppress("UNCHECKED_CAST")
 public fun <PropsT, OutputT, RenderingT> Workflow<PropsT, OutputT, RenderingT>.testRender(
   props: PropsT
 ): RenderTester<PropsT, *, OutputT, RenderingT> {
   val statefulWorkflow = asStatefulWorkflow() as StatefulWorkflow<PropsT, Any?, OutputT, RenderingT>
   return statefulWorkflow.testRender(
     props = props,
-    initialState = statefulWorkflow.initialState(props, null)
+    initialState = run {
+      require(this !is SessionWorkflow<PropsT, *, OutputT, RenderingT>) {
+        "Called testRender on a SessionWorkflow without a CoroutineScope. Use the version that passes a CoroutineScope."
+      }
+      statefulWorkflow.initialState(props, null)
+    }
+  ) as RenderTester<PropsT, Nothing, OutputT, RenderingT>
+}
+
+/**
+ * Create a [RenderTester] to unit test an individual render pass of this [SessionWorkflow],
+ * using the workflow's [initial state][StatefulWorkflow.initialState], in the [workflowScope].
+ *
+ * See [RenderTester] for usage documentation.
+ */
+@OptIn(WorkflowExperimentalApi::class)
+@Suppress("UNCHECKED_CAST")
+public fun <PropsT, OutputT, RenderingT> Workflow<PropsT, OutputT, RenderingT>.testRender(
+  props: PropsT,
+  workflowScope: CoroutineScope
+): RenderTester<PropsT, *, OutputT, RenderingT> {
+  require(this is SessionWorkflow<PropsT, *, OutputT, RenderingT>) {
+    "testRender with workflowScope called on non-Session Workflow. Use the version of testScope with a CoroutineScope to test SessionWorkflow."
+  }
+  val sessionWorkflow: SessionWorkflow<PropsT, Any?, OutputT, RenderingT> =
+    asStatefulWorkflow() as SessionWorkflow<PropsT, Any?, OutputT, RenderingT>
+  return sessionWorkflow.testRender(
+    props = props,
+    initialState = sessionWorkflow.initialState(props, null, workflowScope)
   ) as RenderTester<PropsT, Nothing, OutputT, RenderingT>
 }
 

workflow-testing/src/test/java/com/squareup/workflow1/testing/RealRenderTesterTest.kt

@@ -10,6 +10,7 @@ import com.squareup.workflow1.Worker
 import com.squareup.workflow1.Workflow
 import com.squareup.workflow1.WorkflowAction
 import com.squareup.workflow1.WorkflowAction.Companion.noAction
+import com.squareup.workflow1.WorkflowExperimentalApi
 import com.squareup.workflow1.WorkflowIdentifier
 import com.squareup.workflow1.WorkflowOutput
 import com.squareup.workflow1.action
@@ -19,14 +20,18 @@ import com.squareup.workflow1.identifier
 import com.squareup.workflow1.renderChild
 import com.squareup.workflow1.rendering
 import com.squareup.workflow1.runningWorker
+import com.squareup.workflow1.sessionWorkflow
 import com.squareup.workflow1.stateful
 import com.squareup.workflow1.stateless
 import com.squareup.workflow1.testing.RenderTester.ChildWorkflowMatch.Matched
 import com.squareup.workflow1.unsnapshottableIdentifier
 import com.squareup.workflow1.workflowIdentifier
 import io.mockk.mockk
+import kotlinx.coroutines.CoroutineScope
 import kotlinx.coroutines.flow.Flow
 import kotlinx.coroutines.flow.emptyFlow
+import kotlinx.coroutines.sync.Mutex
+import kotlinx.coroutines.test.runTest
 import org.mockito.kotlin.mock
 import kotlin.reflect.typeOf
 import kotlin.test.Test
@@ -1263,6 +1268,121 @@ internal class RealRenderTesterTest {
     assertEquals(2, renderCount)
   }
 
+  @OptIn(WorkflowExperimentalApi::class)
+  @Test
+  fun `testRender with SessionWorkflow throws exception`() {
+    class TestAction : WorkflowAction<Unit, String, String>() {
+      override fun Updater.apply() {
+        state = "new state"
+        setOutput("output")
+      }
+    }
+
+    val workflow = Workflow.sessionWorkflow<Unit, String, String, Sink<TestAction>>(
+      initialState = { _, _: CoroutineScope -> "initial" },
+      render = { _, _ ->
+        actionSink.contraMap { it }
+      }
+    )
+
+    val exception = assertFailsWith<IllegalArgumentException> {
+      workflow.testRender(Unit)
+        .render { sink ->
+          sink.send(TestAction())
+        }
+    }
+
+    assertEquals(
+      exception.message,
+      "Called testRender on a SessionWorkflow without a CoroutineScope. Use the version that passes a CoroutineScope."
+    )
+  }
+
+  @OptIn(WorkflowExperimentalApi::class)
+  @Test
+  fun `testRender with CoroutineScope works for SessionWorkflow`() = runTest {
+    class TestAction : WorkflowAction<Unit, String, String>() {
+      override fun Updater.apply() {
+        state = "new state"
+        setOutput("output")
+      }
+    }
+
+    val workflow = Workflow.sessionWorkflow<Unit, String, String, Sink<TestAction>>(
+      initialState = { _, _: CoroutineScope -> "initial" },
+      render = { _, _ ->
+        actionSink.contraMap { it }
+      }
+    )
+
+    val testResult = workflow.testRender(Unit, this)
+      .render { sink ->
+        sink.send(TestAction())
+      }
+
+    testResult.verifyActionResult { state, output ->
+      assertEquals("new state", state)
+      assertEquals("output", output?.value)
+    }
+  }
+
+  @OptIn(WorkflowExperimentalApi::class)
+  @Test
+  fun `testRender with CoroutineScope uses the correct scope`() = runTest {
+    val signalMutex = Mutex(locked = true)
+    class TestAction : WorkflowAction<Unit, String, String>() {
+      override fun Updater.apply() {
+        state = "new state"
+        setOutput("output")
+      }
+    }
+
+    val workflow = Workflow.sessionWorkflow<Unit, String, String, Sink<TestAction>>(
+      initialState = { _, workflowScope: CoroutineScope ->
+        assertEquals(workflowScope, this@runTest)
+        signalMutex.unlock()
+        "initial"
+      },
+      render = { _, _ ->
+        actionSink.contraMap { it }
+      }
+    )
+
+    workflow.testRender(Unit, this)
+      .render { sink ->
+        sink.send(TestAction())
+      }
+
+    // Assertion happens in the `initialState` call above.
+    signalMutex.lock()
+  }
+
+  @Test fun `testRender with CoroutineScope does not work if not SessionWorkflow`() = runTest {
+    class TestAction : WorkflowAction<Unit, String, String>() {
+      override fun Updater.apply() {
+        state = "new state"
+        setOutput("output")
+      }
+    }
+
+    val workflow = Workflow.stateful<Unit, String, String, Sink<TestAction>>(
+      initialState = { "initial" },
+      render = { _, _ -> actionSink.contraMap { it } }
+    )
+
+    val exception = assertFailsWith<IllegalArgumentException> {
+      workflow.testRender(Unit, this)
+        .render { sink ->
+          sink.send(TestAction())
+        }
+    }
+
+    assertEquals(
+      exception.message,
+      "testRender with workflowScope called on non-Session Workflow. Use the version of testScope with a CoroutineScope to test SessionWorkflow."
+    )
+  }
+
   @Test fun `createRenderChildInvocation() for Workflow-stateless{}`() {
     val workflow = Workflow.stateless<String, Int, Unit> {}
     val invocation = createRenderChildInvocation(workflow, "props", "key")