Testsuite: Support interceptors in child workflows mocks (#1039)

Co-authored-by: Liang Mei <meiliang@uber.com>

Author: seriousben

internal/internal_workflow.go

@@ -438,7 +438,10 @@ func newWorkflowContext(env workflowEnvironment, interceptors WorkflowIntercepto
 	return rootCtx
 }
 
-func newWorkflowInterceptors(env workflowEnvironment, factories []WorkflowInterceptorFactory) (WorkflowInterceptor, *workflowEnvironmentInterceptor) {
+func newWorkflowInterceptors(
+	env workflowEnvironment,
+	factories []WorkflowInterceptorFactory,
+) (WorkflowInterceptor, *workflowEnvironmentInterceptor) {
 	envInterceptor := &workflowEnvironmentInterceptor{env: env}
 	var interceptor WorkflowInterceptor = envInterceptor
 	for i := len(factories) - 1; i >= 0; i-- {

internal/internal_workflow_testsuite.go

@@ -1394,8 +1394,10 @@ func (w *workflowExecutorWrapper) Execute(ctx Context, input []byte) (result []b
 	// the main loop, but the mock could block if it is configured to wait. So we need to use a separate goroutinue to
 	// run the mock, and resume after mock call returns.
 	mockReadyChannel := NewChannel(ctx)
-	// make a copy of the context for getMockReturn() call to avoid race condition
-	ctxCopy := newWorkflowContext(w.env, nil, nil)
+	// Make a copy of the context for getMockReturn() call to avoid race condition.
+	// Use existing interceptors from env.
+	envInterceptor := &workflowEnvironmentInterceptor{env: env}
+	ctxCopy := newWorkflowContext(w.env, envInterceptor, envInterceptor)
 	go func() {
 		// getMockReturn could block if mock is configured to wait. The returned mockRet is what has been configured
 		// for the mock by using MockCallWrapper.Return(). The mockRet could be mock values or mock function. We process

internal/internal_workflow_testsuite_test.go

@@ -700,6 +700,9 @@ func (s *WorkflowTestSuiteUnitTest) Test_ChildWorkflow_Mock() {
 		if err != nil {
 			return "", err
 		}
+
+		cwo = ChildWorkflowOptions{WorkflowID: "workflow-id", ExecutionStartToCloseTimeout: time.Minute}
+		ctx = WithChildWorkflowOptions(ctx, cwo)
 		var heartbeatWorkflowResult string
 		err = ExecuteChildWorkflow(ctx, testWorkflowHeartbeat, "slow", time.Hour).Get(ctx, &heartbeatWorkflowResult)
 		if err != nil {
@@ -718,6 +721,13 @@ func (s *WorkflowTestSuiteUnitTest) Test_ChildWorkflow_Mock() {
 
 	env.OnActivity(testActivityHello, mock.Anything, mock.Anything).Return("mock_msg", nil)
 	env.OnWorkflow(testWorkflowHeartbeat, mock.Anything, mock.Anything, mock.Anything).
+		Run(func(a mock.Arguments) {
+			ctx := a.Get(0).(Context)
+
+			// Ensure GetWorkflowInfo is usable in mocks.
+			info := GetWorkflowInfo(ctx)
+			s.Equal("workflow-id", info.WorkflowExecution.ID)
+		}).
 		Return("mock_heartbeat", nil)
 	env.ExecuteWorkflow(workflowFn)