From ffa2556f6566b26e53d4a2010127a954def141a7 Mon Sep 17 00:00:00 2001
From: Zach Klippenstein <zach.klippenstein@gmail.com>
Date: Thu, 6 Feb 2020 13:14:39 -0800
Subject: [PATCH] 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.
---
 .../java/com/squareup/workflow/testing/LaunchWorkflow.kt     | 4 +++-
 .../java/com/squareup/workflow/testing/WorkflowTestParams.kt | 5 +++--
 .../java/com/squareup/workflow/testing/WorkflowTester.kt     | 5 ++++-
 3 files changed, 10 insertions(+), 4 deletions(-)

diff --git a/kotlin/workflow-runtime/src/main/java/com/squareup/workflow/testing/LaunchWorkflow.kt b/kotlin/workflow-runtime/src/main/java/com/squareup/workflow/testing/LaunchWorkflow.kt
index 359e1eb15..a034edae2 100644
--- a/kotlin/workflow-runtime/src/main/java/com/squareup/workflow/testing/LaunchWorkflow.kt
+++ b/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
   )
 }
diff --git a/kotlin/workflow-runtime/src/main/java/com/squareup/workflow/testing/WorkflowTestParams.kt b/kotlin/workflow-runtime/src/main/java/com/squareup/workflow/testing/WorkflowTestParams.kt
index c4f5b9b2c..d789287bf 100644
--- a/kotlin/workflow-runtime/src/main/java/com/squareup/workflow/testing/WorkflowTestParams.kt
+++ b/kotlin/workflow-runtime/src/main/java/com/squareup/workflow/testing/WorkflowTestParams.kt
@@ -37,8 +37,9 @@ import kotlin.coroutines.EmptyCoroutineContext
  * 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.
+ * Default is [EmptyCoroutineContext], which means the workers will use their workflow's context
+ * plus the base context passed into the test method. Any context element present in [workerContext]
+ * will override all other context elements for workers.
  */
 @TestOnly
 data class WorkflowTestParams<out StateT>(
diff --git a/kotlin/workflow-testing/src/main/java/com/squareup/workflow/testing/WorkflowTester.kt b/kotlin/workflow-testing/src/main/java/com/squareup/workflow/testing/WorkflowTester.kt
index 7afc74711..168704125 100644
--- a/kotlin/workflow-testing/src/main/java/com/squareup/workflow/testing/WorkflowTester.kt
+++ b/kotlin/workflow-testing/src/main/java/com/squareup/workflow/testing/WorkflowTester.kt
@@ -283,7 +283,10 @@ 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, and that's reasonable behavior.
+      workerContext = context + testParams.workerContext
   ) { session ->
     WorkflowTester(this, propsChannel, session.renderingsAndSnapshots, session.outputs)
         .apply { collectFromWorkflow() }