Make it possible to cancel workflows by ID when testing (#1392)

Make it possible to cancel workflows by ID when testing
temporalio · Feb 21, 2024 · 6d4ec2d · 6d4ec2d
1 parent 5621b7f
commit 6d4ec2d
Show file tree

Hide file tree

Showing 3 changed files with 51 additions and 2 deletions.
diff --git a/internal/internal_workflow_testsuite.go b/internal/internal_workflow_testsuite.go
@@ -2427,12 +2427,16 @@ func (env *testWorkflowEnvironmentImpl) getActivityInfo(activityID ActivityID, a
 }
 
 func (env *testWorkflowEnvironmentImpl) cancelWorkflow(callback ResultHandler) {
+	env.cancelWorkflowByID(env.workflowInfo.WorkflowExecution.ID, env.workflowInfo.WorkflowExecution.RunID, callback)
+}
+
+func (env *testWorkflowEnvironmentImpl) cancelWorkflowByID(workflowID string, runID string, callback ResultHandler) {
 	env.postCallback(func() {
 		// RequestCancelWorkflow needs to be run in main thread
 		env.RequestCancelExternalWorkflow(
 			env.workflowInfo.Namespace,
-			env.workflowInfo.WorkflowExecution.ID,
-			env.workflowInfo.WorkflowExecution.RunID,
+			workflowID,
+			runID,
 			callback,
 		)
 	}, true)

diff --git a/internal/internal_workflow_testsuite_test.go b/internal/internal_workflow_testsuite_test.go
@@ -602,6 +602,46 @@ func (s *WorkflowTestSuiteUnitTest) Test_WorkflowCancellation() {
 	s.True(errors.As(err, &err1))
 }
 
+func (s *WorkflowTestSuiteUnitTest) Test_ChildWorkflowCancellation() {
+	childWorkflowFn := func(ctx Context) error {
+		ctx = WithActivityOptions(ctx, s.activityOptions)
+		f := ExecuteActivity(ctx, testActivityHeartbeat, "msg1", time.Second*10)
+		err := f.Get(ctx, nil) // wait for result
+		return err
+	}
+
+	childID := "child_workflow_id"
+	workflowFn := func(ctx Context) (string, error) {
+		ctx = WithChildWorkflowOptions(ctx, ChildWorkflowOptions{
+			WorkflowID: childID,
+		})
+		err := ExecuteChildWorkflow(ctx, childWorkflowFn).Get(ctx, nil)
+		if err != nil {
+			return err.Error(), nil
+		}
+		return "", nil
+	}
+
+	env := s.NewTestWorkflowEnvironment()
+	env.RegisterDelayedCallback(func() {
+		env.CancelWorkflowByID(childID, "")
+	}, time.Second)
+
+	env.RegisterWorkflow(workflowFn)
+	env.RegisterWorkflow(childWorkflowFn)
+	env.RegisterActivity(testActivityHeartbeat)
+
+	env.ExecuteWorkflow(workflowFn)
+
+	s.True(env.IsWorkflowCompleted())
+	err := env.GetWorkflowError()
+	s.NoError(err)
+
+	var res string
+	s.NoError(env.GetWorkflowResult(&res))
+	s.Contains(res, "canceled")
+}
+
 func testWorkflowHello(ctx Context) (string, error) {
 	ao := ActivityOptions{
 		ScheduleToStartTimeout: time.Minute,

diff --git a/internal/workflow_testsuite.go b/internal/workflow_testsuite.go
@@ -861,6 +861,11 @@ func (e *TestWorkflowEnvironment) CancelWorkflow() {
 	e.impl.cancelWorkflow(func(result *commonpb.Payloads, err error) {})
 }
 
+// CancelWorkflowByID requests cancellation (through workflow Context) to the specified workflow.
+func (e *TestWorkflowEnvironment) CancelWorkflowByID(workflowID string, runID string) {
+	e.impl.cancelWorkflowByID(workflowID, runID, func(result *commonpb.Payloads, err error) {})
+}
+
 // SignalWorkflow sends signal to the currently running test workflow.
 func (e *TestWorkflowEnvironment) SignalWorkflow(name string, input interface{}) {
 	e.impl.signalWorkflow(name, input, true)