Make WorkflowTester tests use the main context as the worker context.

This is an enhancement to #940 that gives more reasonable behavior. The context
parameter to the test methods are primarily used to pass in test dispatchers, and
those dispatchers often really only need to be used for workers.

Since I don't know of a use case for specifying context separately just for workers,
this also removes the `WorkflowTestParams.workerContext` property.

CHANGELOG.md

@@ -5,7 +5,7 @@ Change Log
 
 ### Kotlin
 
- * Make the context used to start workers configurable for tests. (#940)
+ * Make the context used to start workers configurable for tests. (#940, #943)
 
 ### Swift
 

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

@@ -131,6 +131,7 @@ internal fun <PropsT, StateT, OutputT : Any, RenderingT, RunnerT> launchWorkflow
   val renderingsAndSnapshots = ConflatedBroadcastChannel<RenderingAndSnapshot<RenderingT>>()
   val outputs = BroadcastChannel<OutputT>(capacity = 1)
   val workflowScope = scope + Job(parent = scope.coroutineContext[Job])
+  require(workerContext[Job] == null) { "Expected workerContext not to have a Job." }
 
   // Give the caller a chance to start collecting outputs.
   val session = WorkflowSession(renderingsAndSnapshots.asFlow(), outputs.asFlow())

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

@@ -44,6 +44,9 @@ import kotlin.coroutines.EmptyCoroutineContext
  * value to its parent. Returns either the output to be emitted from the root workflow, or null.
  * @param initialState Allows unit tests to start the node from a given state, instead of calling
  * [StatefulWorkflow.initialState].
+ * @param workerContext [CoroutineContext] that is appended to the end of the context used to launch
+ * worker coroutines. This context will override anything from the workflow's scope and any other
+ * hard-coded values added to worker contexts. It must not contain a [Job] element.
  */
 @UseExperimental(VeryExperimentalWorkflow::class)
 internal class WorkflowNode<PropsT, StateT, OutputT : Any, RenderingT>(

kotlin/workflow-runtime/src/main/java/com/squareup/workflow/testing/LaunchWorkflow.kt

@@ -28,6 +28,7 @@ import com.squareup.workflow.testing.WorkflowTestParams.StartMode.StartFromWorkf
 import kotlinx.coroutines.CoroutineScope
 import kotlinx.coroutines.flow.Flow
 import org.jetbrains.annotations.TestOnly
+import kotlin.coroutines.CoroutineContext
 
 /**
  * Launches the [workflow] in a new coroutine in [scope]. The workflow tree is seeded with
@@ -42,6 +43,7 @@ fun <PropsT, StateT, OutputT : Any, RenderingT, RunnerT> launchWorkflowForTestFr
   workflow: StatefulWorkflow<PropsT, StateT, OutputT, RenderingT>,
   props: Flow<PropsT>,
   testParams: WorkflowTestParams<StateT>,
+  workerContext: CoroutineContext,
   beforeStart: CoroutineScope.(session: WorkflowSession<OutputT, RenderingT>) -> RunnerT
 ): RunnerT {
   val initialState = (testParams.startFrom as? StartFromState)?.state
@@ -63,6 +65,6 @@ fun <PropsT, StateT, OutputT : Any, RenderingT, RunnerT> launchWorkflowForTestFr
       initialState = initialState,
       initialSnapshot = initialSnapshot,
       beforeStart = beforeStart,
-      workerContext = testParams.workerContext
+      workerContext = workerContext
   )
 }

kotlin/workflow-runtime/src/main/java/com/squareup/workflow/testing/WorkflowTestParams.kt

@@ -22,8 +22,6 @@ import com.squareup.workflow.testing.WorkflowTestParams.StartMode.StartFromCompl
 import com.squareup.workflow.testing.WorkflowTestParams.StartMode.StartFromState
 import com.squareup.workflow.testing.WorkflowTestParams.StartMode.StartFromWorkflowSnapshot
 import org.jetbrains.annotations.TestOnly
-import kotlin.coroutines.CoroutineContext
-import kotlin.coroutines.EmptyCoroutineContext
 
 /**
  * Defines configuration for workflow testing infrastructure such as `testRender`, `testFromStart`.
@@ -36,15 +34,11 @@ import kotlin.coroutines.EmptyCoroutineContext
  * for any given state, so performing side effects in `render` will almost always result in bugs.
  * It is recommended to leave this on, but if you need to debug a test and don't want to have to
  * deal with the extra passes, you can temporarily set it to false.
- * @param workerContext Used to customize the context in which workers are started for tests.
- * Default is [EmptyCoroutineContext], which means the workers will use the context from their
- * workflow and use the [Unconfined][kotlinx.coroutines.Dispatchers.Unconfined] dispatcher.
  */
 @TestOnly
 data class WorkflowTestParams<out StateT>(
   val startFrom: StartMode<StateT> = StartFresh,
-  val checkRenderIdempotence: Boolean = true,
-  val workerContext: CoroutineContext = EmptyCoroutineContext
+  val checkRenderIdempotence: Boolean = true
 ) {
   /**
    * Defines how to start the workflow for tests.

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

@@ -51,6 +51,7 @@ import kotlin.coroutines.CoroutineContext
 import kotlin.test.AfterTest
 import kotlin.test.assertEquals
 import kotlin.test.assertFails
+import kotlin.test.assertFailsWith
 import kotlin.test.assertFalse
 import kotlin.test.assertNull
 import kotlin.test.assertTrue
@@ -530,6 +531,28 @@ class LaunchWorkflowTest {
       assertEquals(listOf("onRuntimeStopped"), listener.consumeEventNames())
     }
   }
+
+  @Test fun `throws when workerContext contains Job`() {
+    val loop = simpleLoop { _, _ -> hang() }
+    val workflow = Workflow.stateless<Unit, Nothing, Unit> { }
+    val workerContext = Job()
+
+    val error = assertFailsWith<IllegalArgumentException> {
+      runBlocking {
+        launchWorkflowImpl(
+            scope = this,
+            workflowLoop = loop,
+            workflow = workflow.asStatefulWorkflow(),
+            props = emptyFlow(),
+            initialSnapshot = null,
+            initialState = null,
+            beforeStart = {},
+            workerContext = workerContext
+        )
+      }
+    }
+    assertEquals("Expected workerContext not to have a Job.", error.message)
+  }
 }
 
 private suspend fun hang(): Nothing = suspendCancellableCoroutine { }

kotlin/workflow-testing/src/main/java/com/squareup/workflow/testing/WorkflowTester.kt

@@ -27,6 +27,7 @@ import kotlinx.coroutines.CancellationException
 import kotlinx.coroutines.CoroutineExceptionHandler
 import kotlinx.coroutines.CoroutineScope
 import kotlinx.coroutines.Dispatchers.Unconfined
+import kotlinx.coroutines.Job
 import kotlinx.coroutines.NonCancellable
 import kotlinx.coroutines.cancel
 import kotlinx.coroutines.channels.Channel
@@ -283,7 +284,12 @@ fun <T, PropsT, StateT, OutputT : Any, RenderingT>
       scope = CoroutineScope(Unconfined + context + uncaughtExceptionHandler),
       workflow = this@test,
       props = propsChannel.asFlow(),
-      testParams = testParams
+      testParams = testParams,
+      // Use the base context as the starting point for the worker context since it's often used
+      // to pass in test dispatchers.
+      // Remove any job present in the base context since worker context aren't allowed to contain
+      // jobs (it would viol structured concurrency).
+      workerContext = context.minusKey(Job)
   ) { session ->
     WorkflowTester(this, propsChannel, session.renderingsAndSnapshots, session.outputs)
         .apply { collectFromWorkflow() }

kotlin/workflow-testing/src/test/java/com/squareup/workflow/WorkerCompositionIntegrationTest.kt

@@ -19,10 +19,10 @@ package com.squareup.workflow
 
 import com.squareup.workflow.WorkflowAction.Companion.noAction
 import com.squareup.workflow.testing.WorkerSink
-import com.squareup.workflow.testing.WorkflowTestParams
 import com.squareup.workflow.testing.testFromStart
 import kotlinx.coroutines.CancellationException
 import kotlinx.coroutines.CoroutineName
+import kotlinx.coroutines.Job
 import kotlinx.coroutines.TimeoutCancellationException
 import kotlinx.coroutines.channels.Channel
 import kotlin.coroutines.CoroutineContext
@@ -31,6 +31,7 @@ import kotlin.test.Test
 import kotlin.test.assertEquals
 import kotlin.test.assertFailsWith
 import kotlin.test.assertFalse
+import kotlin.test.assertNotSame
 import kotlin.test.assertTrue
 import kotlin.test.fail
 
@@ -250,9 +251,22 @@ class WorkerCompositionIntegrationTest {
     }
     val workerContext = CoroutineName("worker context")
 
-    workflow.testFromStart(testParams = WorkflowTestParams(workerContext = workerContext)) {
+    workflow.testFromStart(context = workerContext) {
       val actualWorkerContext = awaitNextOutput()
       assertEquals("worker context", actualWorkerContext[CoroutineName]!!.name)
     }
   }
+
+  @Test fun `worker context job is ignored`() {
+    val worker = Worker.from { coroutineContext }
+    val workflow = Workflow.stateless<Unit, CoroutineContext, Unit> {
+      runningWorker(worker) { context -> action { setOutput(context) } }
+    }
+    val job: Job = Job()
+
+    workflow.testFromStart(context = job) {
+      val actualWorkerContext = awaitNextOutput()
+      assertNotSame(job, actualWorkerContext[Job])
+    }
+  }
 }