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

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,10 @@ 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 (it would violate
+ * structured concurrency).
  */
 @UseExperimental(VeryExperimentalWorkflow::class)
 internal class WorkflowNode<PropsT, StateT, OutputT : Any, RenderingT>(

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

@@ -26,6 +26,7 @@ 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 kotlinx.coroutines.CoroutineScope
+import kotlinx.coroutines.Job
 import kotlinx.coroutines.flow.Flow
 import org.jetbrains.annotations.TestOnly
 
@@ -63,6 +64,10 @@ fun <PropsT, StateT, OutputT : Any, RenderingT, RunnerT> launchWorkflowForTestFr
       initialState = initialState,
       initialSnapshot = initialSnapshot,
       beforeStart = beforeStart,
-      workerContext = testParams.workerContext
+      // 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 violate structured concurrency).
+      workerContext = scope.coroutineContext.minusKey(Job)
   )
 }

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/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])
+    }
+  }
 }