go.uber.org/cadence@v1.2.9/internal/internal_workflow_testsuite.go (about)

     1  // Copyright (c) 2017-2020 Uber Technologies Inc.
     2  // Portions of the Software are attributed to Copyright (c) 2020 Temporal Technologies Inc.
     3  //
     4  // Permission is hereby granted, free of charge, to any person obtaining a copy
     5  // of this software and associated documentation files (the "Software"), to deal
     6  // in the Software without restriction, including without limitation the rights
     7  // to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
     8  // copies of the Software, and to permit persons to whom the Software is
     9  // furnished to do so, subject to the following conditions:
    10  //
    11  // The above copyright notice and this permission notice shall be included in
    12  // all copies or substantial portions of the Software.
    13  //
    14  // THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
    15  // IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
    16  // FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
    17  // AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
    18  // LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
    19  // OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
    20  // THE SOFTWARE.
    21  
    22  package internal
    23  
    24  import (
    25  	"context"
    26  	"errors"
    27  	"fmt"
    28  	"reflect"
    29  	"strings"
    30  	"sync"
    31  	"time"
    32  
    33  	"github.com/facebookgo/clock"
    34  	"github.com/golang/mock/gomock"
    35  	"github.com/opentracing/opentracing-go"
    36  	"github.com/robfig/cron"
    37  	"github.com/stretchr/testify/mock"
    38  	"github.com/uber-go/tally"
    39  	"go.uber.org/yarpc"
    40  	"go.uber.org/zap"
    41  
    42  	"go.uber.org/cadence/.gen/go/cadence/workflowserviceclient"
    43  	"go.uber.org/cadence/.gen/go/cadence/workflowservicetest"
    44  	"go.uber.org/cadence/.gen/go/shared"
    45  	"go.uber.org/cadence/internal/common"
    46  	"go.uber.org/cadence/internal/common/metrics"
    47  )
    48  
    49  const (
    50  	defaultTestDomain           = "default-test-domain"
    51  	defaultTestTaskList         = "default-test-tasklist"
    52  	defaultTestWorkflowID       = "default-test-workflow-id"
    53  	defaultTestRunID            = "default-test-run-id"
    54  	defaultTestWorkflowTypeName = "default-test-workflow-type-name"
    55  	defaultTestDomainName       = "default-test-domain-name"
    56  	workflowTypeNotSpecified    = "workflow-type-not-specified"
    57  )
    58  
    59  type (
    60  	testTimerHandle struct {
    61  		env            *testWorkflowEnvironmentImpl
    62  		callback       resultHandler
    63  		timer          *clock.Timer
    64  		wallTimer      *clock.Timer
    65  		duration       time.Duration
    66  		mockTimeToFire time.Time
    67  		wallTimeToFire time.Time
    68  		timerID        int
    69  	}
    70  
    71  	testActivityHandle struct {
    72  		callback         resultHandler
    73  		activityType     string
    74  		heartbeatDetails []byte
    75  	}
    76  
    77  	testWorkflowHandle struct {
    78  		env      *testWorkflowEnvironmentImpl
    79  		callback resultHandler
    80  		handled  bool
    81  		params   *executeWorkflowParams
    82  		err      error
    83  	}
    84  
    85  	testCallbackHandle struct {
    86  		callback          func()
    87  		startDecisionTask bool // start a new decision task after callback() is handled.
    88  		env               *testWorkflowEnvironmentImpl
    89  	}
    90  
    91  	activityExecutorWrapper struct {
    92  		*activityExecutor
    93  		env *testWorkflowEnvironmentImpl
    94  	}
    95  
    96  	workflowExecutorWrapper struct {
    97  		*workflowExecutor
    98  		env *testWorkflowEnvironmentImpl
    99  	}
   100  
   101  	mockWrapper struct {
   102  		env           *testWorkflowEnvironmentImpl
   103  		name          string
   104  		fn            interface{}
   105  		isWorkflow    bool
   106  		dataConverter DataConverter
   107  	}
   108  
   109  	taskListSpecificActivity struct {
   110  		fn        interface{}
   111  		taskLists map[string]struct{}
   112  	}
   113  
   114  	// testWorkflowEnvironmentShared is the shared data between parent workflow and child workflow test environments
   115  	testWorkflowEnvironmentShared struct {
   116  		locker    sync.Mutex
   117  		testSuite *WorkflowTestSuite
   118  
   119  		taskListSpecificActivities map[string]*taskListSpecificActivity
   120  
   121  		mock         *mock.Mock
   122  		service      workflowserviceclient.Interface
   123  		logger       *zap.Logger
   124  		metricsScope *metrics.TaggedScope
   125  		ctxProps     []ContextPropagator
   126  		mockClock    *clock.Mock
   127  		wallClock    clock.Clock
   128  		startTime    time.Time
   129  
   130  		callbackChannel chan testCallbackHandle
   131  		testTimeout     time.Duration
   132  		header          *shared.Header
   133  
   134  		counterID        int
   135  		activities       map[string]*testActivityHandle
   136  		localActivities  map[string]*localActivityTask
   137  		timers           map[string]*testTimerHandle
   138  		runningWorkflows map[string]*testWorkflowHandle
   139  
   140  		runningCount int
   141  
   142  		expectedMockCalls map[string]struct{}
   143  
   144  		onActivityStartedListener        func(activityInfo *ActivityInfo, ctx context.Context, args Values)
   145  		onActivityCompletedListener      func(activityInfo *ActivityInfo, result Value, err error)
   146  		onActivityCanceledListener       func(activityInfo *ActivityInfo)
   147  		onLocalActivityStartedListener   func(activityInfo *ActivityInfo, ctx context.Context, args []interface{})
   148  		onLocalActivityCompletedListener func(activityInfo *ActivityInfo, result Value, err error)
   149  		onLocalActivityCanceledListener  func(activityInfo *ActivityInfo)
   150  		onActivityHeartbeatListener      func(activityInfo *ActivityInfo, details Values)
   151  		onChildWorkflowStartedListener   func(workflowInfo *WorkflowInfo, ctx Context, args Values)
   152  		onChildWorkflowCompletedListener func(workflowInfo *WorkflowInfo, result Value, err error)
   153  		onChildWorkflowCanceledListener  func(workflowInfo *WorkflowInfo)
   154  		onTimerScheduledListener         func(timerID string, duration time.Duration)
   155  		onTimerFiredListener             func(timerID string)
   156  		onTimerCancelledListener         func(timerID string)
   157  
   158  		cronMaxIterations int
   159  	}
   160  
   161  	// testWorkflowEnvironmentImpl is the environment that runs the workflow/activity unit tests.
   162  	testWorkflowEnvironmentImpl struct {
   163  		*testWorkflowEnvironmentShared
   164  		parentEnv            *testWorkflowEnvironmentImpl
   165  		registry             *registry
   166  		workflowInterceptors []WorkflowInterceptorFactory
   167  
   168  		workflowInfo   *WorkflowInfo
   169  		workflowDef    workflowDefinition
   170  		changeVersions map[string]Version
   171  		openSessions   map[string]*SessionInfo
   172  
   173  		workflowCancelHandler func()
   174  		signalHandler         func(name string, input []byte)
   175  		queryHandler          func(string, []byte) ([]byte, error)
   176  		startedHandler        func(r WorkflowExecution, e error)
   177  
   178  		isTestCompleted  bool
   179  		testResult       Value
   180  		testError        error
   181  		doneChannel      chan struct{}
   182  		workerOptions    WorkerOptions
   183  		executionTimeout time.Duration
   184  
   185  		heartbeatDetails []byte
   186  
   187  		workerStopChannel  chan struct{}
   188  		sessionEnvironment *testSessionEnvironmentImpl
   189  
   190  		cronSchedule   string
   191  		cronIterations int
   192  		workflowInput  []byte
   193  	}
   194  
   195  	testSessionEnvironmentImpl struct {
   196  		*sessionEnvironmentImpl
   197  		testWorkflowEnvironment *testWorkflowEnvironmentImpl
   198  	}
   199  )
   200  
   201  // make sure interface is implemented
   202  var _ workflowEnvironment = (*testWorkflowEnvironmentImpl)(nil)
   203  
   204  func newTestWorkflowEnvironmentImpl(s *WorkflowTestSuite, parentRegistry *registry) *testWorkflowEnvironmentImpl {
   205  	var r *registry
   206  	if parentRegistry == nil {
   207  		r = newRegistry()
   208  		r.RegisterActivityWithOptions(sessionCreationActivity, RegisterActivityOptions{
   209  			Name: sessionCreationActivityName,
   210  		})
   211  		r.RegisterActivityWithOptions(sessionCompletionActivity, RegisterActivityOptions{
   212  			Name: sessionCompletionActivityName,
   213  		})
   214  	} else {
   215  		r = parentRegistry
   216  	}
   217  
   218  	env := &testWorkflowEnvironmentImpl{
   219  		testWorkflowEnvironmentShared: &testWorkflowEnvironmentShared{
   220  			testSuite:                  s,
   221  			taskListSpecificActivities: make(map[string]*taskListSpecificActivity),
   222  
   223  			logger:           s.logger,
   224  			metricsScope:     metrics.NewTaggedScope(s.scope),
   225  			mockClock:        clock.NewMock(),
   226  			wallClock:        clock.New(),
   227  			timers:           make(map[string]*testTimerHandle),
   228  			activities:       make(map[string]*testActivityHandle),
   229  			localActivities:  make(map[string]*localActivityTask),
   230  			runningWorkflows: make(map[string]*testWorkflowHandle),
   231  			callbackChannel:  make(chan testCallbackHandle, 1000),
   232  			testTimeout:      time.Second * 3,
   233  
   234  			expectedMockCalls: make(map[string]struct{}),
   235  
   236  			cronMaxIterations: -1,
   237  		},
   238  
   239  		workflowInfo: &WorkflowInfo{
   240  			Domain: defaultTestDomain,
   241  			WorkflowExecution: WorkflowExecution{
   242  				ID:    defaultTestWorkflowID,
   243  				RunID: defaultTestRunID,
   244  			},
   245  			WorkflowType: WorkflowType{Name: workflowTypeNotSpecified},
   246  			TaskListName: defaultTestTaskList,
   247  
   248  			ExecutionStartToCloseTimeoutSeconds: 1,
   249  			TaskStartToCloseTimeoutSeconds:      1,
   250  		},
   251  		registry: r,
   252  
   253  		changeVersions: make(map[string]Version),
   254  		openSessions:   make(map[string]*SessionInfo),
   255  
   256  		doneChannel:       make(chan struct{}),
   257  		workerStopChannel: make(chan struct{}),
   258  
   259  		cronIterations: 0,
   260  	}
   261  
   262  	// move forward the mock clock to start time.
   263  	env.setStartTime(time.Now())
   264  
   265  	// put current workflow as a running workflow so child can send signal to parent
   266  	testWorkflowHandle := &testWorkflowHandle{env: env, callback: func(result []byte, err error) {}}
   267  	if env.workflowInfo.CronSchedule != nil && len(*env.workflowInfo.CronSchedule) > 0 {
   268  		testWorkflowHandle.params.cronSchedule = *env.workflowInfo.CronSchedule
   269  	}
   270  	env.runningWorkflows[env.workflowInfo.WorkflowExecution.ID] = testWorkflowHandle
   271  
   272  	if env.logger == nil {
   273  		logger, _ := zap.NewDevelopment()
   274  		env.logger = logger
   275  	}
   276  	if env.metricsScope == nil {
   277  		env.metricsScope = metrics.NewTaggedScope(s.scope)
   278  	}
   279  	env.ctxProps = s.ctxProps
   280  	env.header = s.header
   281  
   282  	// setup mock service
   283  	mockCtrl := gomock.NewController(&testReporter{logger: env.logger})
   284  	mockService := workflowservicetest.NewMockClient(mockCtrl)
   285  
   286  	mockHeartbeatFn := func(c context.Context, r *shared.RecordActivityTaskHeartbeatRequest, opts ...yarpc.CallOption) error {
   287  		activityID := string(r.TaskToken)
   288  		env.locker.Lock() // need lock as this is running in activity worker's goroutinue
   289  		activityHandle, ok := env.getActivityHandle(activityID)
   290  		env.locker.Unlock()
   291  		if !ok {
   292  			env.logger.Debug("RecordActivityTaskHeartbeat: ActivityID not found, could be already completed or cancelled.",
   293  				zap.String(tagActivityID, activityID))
   294  			return &shared.EntityNotExistsError{}
   295  		}
   296  		activityHandle.heartbeatDetails = r.Details
   297  		activityInfo := env.getActivityInfo(activityID, activityHandle.activityType)
   298  		env.postCallback(func() {
   299  			if env.onActivityHeartbeatListener != nil {
   300  				env.onActivityHeartbeatListener(activityInfo, newEncodedValues(r.Details, env.GetDataConverter()))
   301  			}
   302  		}, false)
   303  
   304  		env.logger.Debug("RecordActivityTaskHeartbeat", zap.String(tagActivityID, activityID))
   305  		return nil
   306  	}
   307  
   308  	var callOptions []interface{}
   309  	yarpcCallOptions := getYarpcCallOptions(FeatureFlags{})
   310  	for range yarpcCallOptions {
   311  		callOptions = append(callOptions, gomock.Any())
   312  	}
   313  	em := mockService.EXPECT().RecordActivityTaskHeartbeat(gomock.Any(), gomock.Any(), callOptions...).
   314  		Return(&shared.RecordActivityTaskHeartbeatResponse{CancelRequested: common.BoolPtr(false)}, nil)
   315  	em.Do(func(ctx context.Context, r *shared.RecordActivityTaskHeartbeatRequest, opts ...yarpc.CallOption) {
   316  		// TODO: The following will hit a data race in the gomock code where the Do() action is executed outside
   317  		// the lock and setting return value from inside the action is going to run into races.
   318  		// err := mockHeartbeatFn(ctx, r, opts)
   319  		// em.Return(&shared.RecordActivityTaskHeartbeatResponse{CancelRequested: common.BoolPtr(false)}, err)
   320  		mockHeartbeatFn(ctx, r, opts...)
   321  	}).AnyTimes()
   322  
   323  	env.service = mockService
   324  
   325  	if env.workerOptions.Logger == nil {
   326  		env.workerOptions.Logger = env.logger
   327  	}
   328  	if env.workerOptions.MetricsScope == nil {
   329  		env.workerOptions.MetricsScope = env.metricsScope
   330  	}
   331  	if env.workerOptions.DataConverter == nil {
   332  		env.workerOptions.DataConverter = getDefaultDataConverter()
   333  	}
   334  	if len(env.workerOptions.ContextPropagators) == 0 {
   335  		env.workerOptions.ContextPropagators = env.ctxProps
   336  	}
   337  
   338  	return env
   339  }
   340  
   341  func (env *testWorkflowEnvironmentImpl) setStartTime(startTime time.Time) {
   342  	// move forward the mock clock to start time.
   343  	if startTime.IsZero() {
   344  		// if start time not set, use current clock time
   345  		startTime = env.wallClock.Now()
   346  	}
   347  	env.mockClock.Add(startTime.Sub(env.mockClock.Now()))
   348  }
   349  
   350  func (env *testWorkflowEnvironmentImpl) setCronSchedule(cronSchedule string) {
   351  	env.workflowInfo.CronSchedule = &cronSchedule
   352  }
   353  
   354  func (env *testWorkflowEnvironmentImpl) setCronMaxIterationas(cronMaxIterations int) {
   355  	env.cronMaxIterations = cronMaxIterations
   356  }
   357  
   358  func (env *testWorkflowEnvironmentImpl) newTestWorkflowEnvironmentForChild(params *executeWorkflowParams, callback resultHandler, startedHandler func(r WorkflowExecution, e error)) (*testWorkflowEnvironmentImpl, error) {
   359  	// create a new test env
   360  	childEnv := newTestWorkflowEnvironmentImpl(env.testSuite, env.registry)
   361  	childEnv.parentEnv = env
   362  	childEnv.startedHandler = startedHandler
   363  	childEnv.testWorkflowEnvironmentShared = env.testWorkflowEnvironmentShared
   364  	childEnv.workerOptions = env.workerOptions
   365  	childEnv.workerOptions.DataConverter = params.dataConverter
   366  	childEnv.workflowInterceptors = env.workflowInterceptors
   367  	childEnv.registry = env.registry
   368  
   369  	if params.workflowID == "" {
   370  		params.workflowID = env.workflowInfo.WorkflowExecution.RunID + "_" + getStringID(env.nextID())
   371  	}
   372  	var cronSchedule *string
   373  	if len(params.cronSchedule) > 0 {
   374  		cronSchedule = &params.cronSchedule
   375  	}
   376  	// set workflow info data for child workflow
   377  	childEnv.header = params.header
   378  	childEnv.workflowInfo.Attempt = params.attempt
   379  	childEnv.workflowInfo.WorkflowExecution.ID = params.workflowID
   380  	childEnv.workflowInfo.WorkflowExecution.RunID = params.workflowID + "_RunID"
   381  	childEnv.workflowInfo.Domain = *params.domain
   382  	childEnv.workflowInfo.TaskListName = *params.taskListName
   383  	childEnv.workflowInfo.ExecutionStartToCloseTimeoutSeconds = *params.executionStartToCloseTimeoutSeconds
   384  	childEnv.workflowInfo.TaskStartToCloseTimeoutSeconds = *params.taskStartToCloseTimeoutSeconds
   385  	childEnv.workflowInfo.lastCompletionResult = params.lastCompletionResult
   386  	childEnv.workflowInfo.CronSchedule = cronSchedule
   387  	childEnv.workflowInfo.ParentWorkflowDomain = &env.workflowInfo.Domain
   388  	childEnv.workflowInfo.ParentWorkflowExecution = &env.workflowInfo.WorkflowExecution
   389  	childEnv.executionTimeout = time.Duration(*params.executionStartToCloseTimeoutSeconds) * time.Second
   390  	if workflowHandler, ok := env.runningWorkflows[params.workflowID]; ok {
   391  		// duplicate workflow ID
   392  		if !workflowHandler.handled {
   393  			return nil, &shared.WorkflowExecutionAlreadyStartedError{
   394  				Message: common.StringPtr("Workflow execution already started"),
   395  			}
   396  		}
   397  		if params.workflowIDReusePolicy == WorkflowIDReusePolicyRejectDuplicate {
   398  			return nil, &shared.WorkflowExecutionAlreadyStartedError{
   399  				Message: common.StringPtr("Workflow execution already started"),
   400  			}
   401  		}
   402  		if workflowHandler.err == nil && params.workflowIDReusePolicy == WorkflowIDReusePolicyAllowDuplicateFailedOnly {
   403  			return nil, &shared.WorkflowExecutionAlreadyStartedError{
   404  				Message: common.StringPtr("Workflow execution already started"),
   405  			}
   406  		}
   407  	}
   408  
   409  	env.runningWorkflows[params.workflowID] = &testWorkflowHandle{env: childEnv, callback: callback, params: params}
   410  
   411  	return childEnv, nil
   412  }
   413  
   414  func (env *testWorkflowEnvironmentImpl) setWorkerOptions(options WorkerOptions) {
   415  	if len(options.Identity) > 0 {
   416  		env.workerOptions.Identity = options.Identity
   417  	}
   418  	if options.BackgroundActivityContext != nil {
   419  		env.workerOptions.BackgroundActivityContext = options.BackgroundActivityContext
   420  	}
   421  	if options.MetricsScope != nil {
   422  		env.workerOptions.MetricsScope = options.MetricsScope
   423  	}
   424  	if options.DataConverter != nil {
   425  		env.workerOptions.DataConverter = options.DataConverter
   426  	}
   427  	// Uncomment when resourceID is exposed to user.
   428  	// if options.SessionResourceID != "" {
   429  	// 	env.workerOptions.SessionResourceID = options.SessionResourceID
   430  	// }
   431  	if options.MaxConcurrentSessionExecutionSize != 0 {
   432  		env.workerOptions.MaxConcurrentSessionExecutionSize = options.MaxConcurrentSessionExecutionSize
   433  	}
   434  	if len(options.ContextPropagators) > 0 {
   435  		env.workerOptions.ContextPropagators = options.ContextPropagators
   436  	}
   437  	if options.Logger != nil {
   438  		env.workerOptions.Logger = options.Logger
   439  	}
   440  	env.workflowInterceptors = options.WorkflowInterceptorChainFactories
   441  }
   442  
   443  func (env *testWorkflowEnvironmentImpl) setWorkerStopChannel(c chan struct{}) {
   444  	env.workerStopChannel = c
   445  }
   446  
   447  func (env *testWorkflowEnvironmentImpl) setActivityTaskList(tasklist string, activityFns ...interface{}) {
   448  	for _, activityFn := range activityFns {
   449  		fnName := getActivityFunctionName(env.registry, activityFn)
   450  		taskListActivity, ok := env.taskListSpecificActivities[fnName]
   451  		if !ok {
   452  			taskListActivity = &taskListSpecificActivity{fn: activityFn, taskLists: make(map[string]struct{})}
   453  			env.taskListSpecificActivities[fnName] = taskListActivity
   454  		}
   455  		taskListActivity.taskLists[tasklist] = struct{}{}
   456  	}
   457  }
   458  
   459  func (env *testWorkflowEnvironmentImpl) executeWorkflow(workflowFn interface{}, args ...interface{}) {
   460  	fType := reflect.TypeOf(workflowFn)
   461  	if getKind(fType) == reflect.Func {
   462  		env.RegisterWorkflowWithOptions(workflowFn, RegisterWorkflowOptions{DisableAlreadyRegisteredCheck: true})
   463  	}
   464  	workflowType, input, err := getValidatedWorkflowFunction(workflowFn, args, env.GetDataConverter(), env.GetRegistry())
   465  	if err != nil {
   466  		panic(err)
   467  	}
   468  	env.executeWorkflowInternal(0, workflowType.Name, input)
   469  }
   470  
   471  func (env *testWorkflowEnvironmentImpl) executeWorkflowInternal(delayStart time.Duration, workflowType string, input []byte) {
   472  	env.locker.Lock()
   473  	if env.workflowInfo.WorkflowType.Name != workflowTypeNotSpecified {
   474  		// Current TestWorkflowEnvironment only support to run one workflow.
   475  		// Created task to support testing multiple workflows with one env instance
   476  		// https://github.com/uber-go/cadence-client/issues/616
   477  		panic(fmt.Sprintf("Current TestWorkflowEnvironment is used to execute %v. Please create a new TestWorkflowEnvironment for %v.", env.workflowInfo.WorkflowType.Name, workflowType))
   478  	}
   479  	env.workflowInfo.WorkflowType.Name = workflowType
   480  	env.locker.Unlock()
   481  
   482  	workflowDefinition, err := env.getWorkflowDefinition(env.workflowInfo.WorkflowType)
   483  	if err != nil {
   484  		panic(err)
   485  	}
   486  	env.workflowDef = workflowDefinition
   487  	// Store the Workflow input for potential Cron
   488  	env.workflowInput = input
   489  
   490  	// env.workflowDef.Execute() method will execute dispatcher. We want the dispatcher to only run in main loop.
   491  	// In case of child workflow, this executeWorkflowInternal() is run in separate goroutinue, so use postCallback
   492  	// to make sure workflowDef.Execute() is run in main loop.
   493  	env.postCallback(func() {
   494  		env.workflowDef.Execute(env, env.header, input)
   495  		// kick off first decision task to start the workflow
   496  		if delayStart == 0 {
   497  			env.startDecisionTask()
   498  		} else {
   499  			// we need to delayStart start workflow, decrease runningCount so mockClock could auto forward
   500  			env.runningCount--
   501  			env.registerDelayedCallback(func() {
   502  				env.runningCount++
   503  				env.startDecisionTask()
   504  			}, delayStart)
   505  		}
   506  	}, false)
   507  
   508  	if env.executionTimeout > 0 {
   509  		timeoutDuration := env.executionTimeout + delayStart
   510  		env.registerDelayedCallback(func() {
   511  			if !env.isTestCompleted {
   512  				env.Complete(nil, ErrDeadlineExceeded)
   513  			}
   514  		}, timeoutDuration)
   515  	}
   516  	env.startMainLoop()
   517  }
   518  
   519  func (env *testWorkflowEnvironmentImpl) getWorkflowDefinition(wt WorkflowType) (workflowDefinition, error) {
   520  	wf, ok := env.registry.getWorkflowFn(wt.Name)
   521  	if !ok {
   522  		supported := strings.Join(env.registry.GetRegisteredWorkflowTypes(), ", ")
   523  		return nil, fmt.Errorf("unable to find workflow type: %v. Supported types: [%v]", wt.Name, supported)
   524  	}
   525  	wd := &workflowExecutorWrapper{
   526  		workflowExecutor: &workflowExecutor{workflowType: wt.Name, fn: wf},
   527  		env:              env,
   528  	}
   529  	return newSyncWorkflowDefinition(wd), nil
   530  }
   531  
   532  func (env *testWorkflowEnvironmentImpl) executeActivity(
   533  	activityFn interface{},
   534  	args ...interface{},
   535  ) (Value, error) {
   536  	return env.executeActivityWithOptions(
   537  		activityOptions{
   538  			ScheduleToCloseTimeoutSeconds: 600,
   539  			StartToCloseTimeoutSeconds:    600,
   540  		},
   541  		activityFn,
   542  		args...,
   543  	)
   544  }
   545  
   546  func (env *testWorkflowEnvironmentImpl) executeActivityWithOptions(
   547  	activityOptions activityOptions,
   548  	activityFn interface{},
   549  	args ...interface{},
   550  ) (Value, error) {
   551  	activityType, err := getValidatedActivityFunction(activityFn, args, env.registry)
   552  	if err != nil {
   553  		panic(err)
   554  	}
   555  
   556  	input, err := encodeArgs(env.GetDataConverter(), args)
   557  	if err != nil {
   558  		panic(err)
   559  	}
   560  
   561  	params := executeActivityParams{
   562  		activityOptions: activityOptions,
   563  		ActivityType:    *activityType,
   564  		Input:           input,
   565  		Header:          env.header,
   566  	}
   567  
   568  	task := newTestActivityTask(
   569  		defaultTestWorkflowID,
   570  		defaultTestRunID,
   571  		"0",
   572  		defaultTestWorkflowTypeName,
   573  		defaultTestDomainName,
   574  		params,
   575  	)
   576  
   577  	task.HeartbeatDetails = env.heartbeatDetails
   578  
   579  	// ensure activityFn is registered to defaultTestTaskList
   580  	taskHandler := env.newTestActivityTaskHandler(defaultTestTaskList, env.GetDataConverter())
   581  	result, err := taskHandler.Execute(defaultTestTaskList, task)
   582  	if err != nil {
   583  		if err == context.DeadlineExceeded {
   584  			env.logger.Debug(fmt.Sprintf("Activity %v timed out", task.ActivityType.Name))
   585  			return nil, NewTimeoutError(shared.TimeoutTypeStartToClose, context.DeadlineExceeded.Error())
   586  		}
   587  		topLine := fmt.Sprintf("activity for %s [panic]:", defaultTestTaskList)
   588  		st := getStackTraceRaw(topLine, 7, 0)
   589  		return nil, newPanicError(err.Error(), st)
   590  	}
   591  
   592  	if result == ErrActivityResultPending {
   593  		return nil, ErrActivityResultPending
   594  	}
   595  
   596  	switch request := result.(type) {
   597  	case *shared.RespondActivityTaskCanceledRequest:
   598  		details := newEncodedValues(request.Details, env.GetDataConverter())
   599  		return nil, NewCanceledError(details)
   600  	case *shared.RespondActivityTaskFailedRequest:
   601  		return nil, constructError(request.GetReason(), request.Details, env.GetDataConverter())
   602  	case *shared.RespondActivityTaskCompletedRequest:
   603  		return newEncodedValue(request.Result, env.GetDataConverter()), nil
   604  	default:
   605  		// will never happen
   606  		return nil, fmt.Errorf("unsupported respond type %T", result)
   607  	}
   608  }
   609  
   610  func (env *testWorkflowEnvironmentImpl) executeLocalActivity(
   611  	activityFn interface{},
   612  	args ...interface{},
   613  ) (val Value, err error) {
   614  	params := executeLocalActivityParams{
   615  		localActivityOptions: localActivityOptions{
   616  			ScheduleToCloseTimeoutSeconds: common.Int32Ceil(env.testTimeout.Seconds()),
   617  		},
   618  		ActivityFn:   activityFn,
   619  		InputArgs:    args,
   620  		WorkflowInfo: env.workflowInfo,
   621  	}
   622  	task := &localActivityTask{
   623  		activityID: "test-local-activity",
   624  		params:     &params,
   625  		callback: func(lar *localActivityResultWrapper) {
   626  		},
   627  	}
   628  	taskHandler := localActivityTaskHandler{
   629  		userContext:  env.workerOptions.BackgroundActivityContext,
   630  		metricsScope: env.metricsScope,
   631  		logger:       env.logger,
   632  		tracer:       opentracing.NoopTracer{},
   633  	}
   634  
   635  	result := taskHandler.executeLocalActivityTask(task)
   636  	if result.err != nil {
   637  		return nil, result.err
   638  	}
   639  	return newEncodedValue(result.result, env.GetDataConverter()), nil
   640  }
   641  
   642  func (env *testWorkflowEnvironmentImpl) startDecisionTask() {
   643  	if !env.isTestCompleted {
   644  		env.workflowDef.OnDecisionTaskStarted()
   645  	}
   646  }
   647  
   648  func (env *testWorkflowEnvironmentImpl) isChildWorkflow() bool {
   649  	return env.parentEnv != nil
   650  }
   651  
   652  func (env *testWorkflowEnvironmentImpl) startMainLoop() {
   653  	if env.isChildWorkflow() {
   654  		// child workflow rely on parent workflow's main loop to process events
   655  		<-env.doneChannel // wait until workflow is complete
   656  		return
   657  	}
   658  
   659  	for !env.isTestCompleted {
   660  		// use non-blocking-select to check if there is anything pending in the main thread.
   661  		select {
   662  		case c := <-env.callbackChannel:
   663  			// this will drain the callbackChannel
   664  			c.processCallback()
   665  		default:
   666  			// nothing to process, main thread is blocked at this moment, now check if we should auto fire next timer
   667  			if !env.autoFireNextTimer() {
   668  				if env.isTestCompleted {
   669  					return
   670  				}
   671  
   672  				// no timer to fire, wait for things to do or timeout.
   673  				select {
   674  				case c := <-env.callbackChannel:
   675  					c.processCallback()
   676  				case <-time.After(env.testTimeout):
   677  					// not able to complete workflow within test timeout, workflow likely stuck somewhere,
   678  					// check workflow stack for more details.
   679  					panicMsg := fmt.Sprintf("test timeout: %v, workflow stack: %v",
   680  						env.testTimeout, env.workflowDef.StackTrace())
   681  					panic(panicMsg)
   682  				}
   683  			}
   684  		}
   685  	}
   686  }
   687  
   688  func (env *testWorkflowEnvironmentImpl) registerDelayedCallback(f func(), delayDuration time.Duration) {
   689  	timerCallback := func(result []byte, err error) {
   690  		f()
   691  	}
   692  	if delayDuration == 0 {
   693  		env.postCallback(f, false)
   694  		return
   695  	}
   696  	mainLoopCallback := func() {
   697  		env.newTimer(delayDuration, timerCallback, false)
   698  	}
   699  	env.postCallback(mainLoopCallback, false)
   700  }
   701  
   702  func (c *testCallbackHandle) processCallback() {
   703  	c.env.locker.Lock()
   704  	defer c.env.locker.Unlock()
   705  	c.callback()
   706  	if c.startDecisionTask {
   707  		c.env.startDecisionTask()
   708  	}
   709  }
   710  
   711  func (env *testWorkflowEnvironmentImpl) autoFireNextTimer() bool {
   712  	if len(env.timers) == 0 {
   713  		return false
   714  	}
   715  
   716  	// find next timer
   717  	var nextTimer *testTimerHandle
   718  	for _, t := range env.timers {
   719  		if nextTimer == nil {
   720  			nextTimer = t
   721  		} else if t.mockTimeToFire.Before(nextTimer.mockTimeToFire) ||
   722  			(t.mockTimeToFire.Equal(nextTimer.mockTimeToFire) && t.timerID < nextTimer.timerID) {
   723  			nextTimer = t
   724  		}
   725  	}
   726  
   727  	// function to fire timer
   728  	fireTimer := func(th *testTimerHandle) {
   729  		skipDuration := th.mockTimeToFire.Sub(env.mockClock.Now())
   730  		env.logger.Debug("Auto fire timer",
   731  			zap.Int(tagTimerID, th.timerID),
   732  			zap.Duration("TimerDuration", th.duration),
   733  			zap.Duration("TimeSkipped", skipDuration))
   734  
   735  		// Move mockClock forward, this will fire the timer, and the timer callback will remove timer from timers.
   736  		env.mockClock.Add(skipDuration)
   737  	}
   738  
   739  	// fire timer if there is no running activity
   740  	if env.runningCount == 0 {
   741  		if nextTimer.wallTimer != nil {
   742  			nextTimer.wallTimer.Stop()
   743  			nextTimer.wallTimer = nil
   744  		}
   745  		fireTimer(nextTimer)
   746  		return true
   747  	}
   748  
   749  	durationToFire := nextTimer.mockTimeToFire.Sub(env.mockClock.Now())
   750  	wallTimeToFire := env.wallClock.Now().Add(durationToFire)
   751  
   752  	if nextTimer.wallTimer != nil && nextTimer.wallTimeToFire.Before(wallTimeToFire) {
   753  		// nextTimer already set, meaning we already have a wall clock timer for the nextTimer setup earlier. And the
   754  		// previously scheduled wall time to fire is before the wallTimeToFire calculated this time. This could happen
   755  		// if workflow was blocked while there was activity running, and when that activity completed, there are some
   756  		// other activities still running while the nextTimer is still that same nextTimer. In that case, we should not
   757  		// reset the wall time to fire for the nextTimer.
   758  		return false
   759  	}
   760  	if nextTimer.wallTimer != nil {
   761  		// wallTimer was scheduled, but the wall time to fire should be earlier based on current calculation.
   762  		nextTimer.wallTimer.Stop()
   763  	}
   764  
   765  	// there is running activities, we would fire next timer only if wall time passed by nextTimer duration.
   766  	nextTimer.wallTimeToFire, nextTimer.wallTimer = wallTimeToFire, env.wallClock.AfterFunc(durationToFire, func() {
   767  		// make sure it is running in the main loop
   768  		nextTimer.env.postCallback(func() {
   769  			if timerHandle, ok := env.timers[getStringID(nextTimer.timerID)]; ok {
   770  				fireTimer(timerHandle)
   771  			}
   772  		}, true)
   773  	})
   774  
   775  	return false
   776  }
   777  
   778  func (env *testWorkflowEnvironmentImpl) postCallback(cb func(), startDecisionTask bool) {
   779  	env.callbackChannel <- testCallbackHandle{callback: cb, startDecisionTask: startDecisionTask, env: env}
   780  }
   781  
   782  func (env *testWorkflowEnvironmentImpl) RequestCancelActivity(activityID string) {
   783  	handle, ok := env.getActivityHandle(activityID)
   784  	if !ok {
   785  		env.logger.Debug("RequestCancelActivity failed, Activity not exists or already completed.", zap.String(tagActivityID, activityID))
   786  		return
   787  	}
   788  	activityInfo := env.getActivityInfo(activityID, handle.activityType)
   789  	env.logger.Debug("RequestCancelActivity", zap.String(tagActivityID, activityID))
   790  	env.deleteHandle(activityID)
   791  	env.postCallback(func() {
   792  		handle.callback(nil, NewCanceledError())
   793  		if env.onActivityCanceledListener != nil {
   794  			env.onActivityCanceledListener(activityInfo)
   795  		}
   796  	}, true)
   797  }
   798  
   799  // RequestCancelTimer request to cancel timer on this testWorkflowEnvironmentImpl.
   800  func (env *testWorkflowEnvironmentImpl) RequestCancelTimer(timerID string) {
   801  	env.logger.Debug("RequestCancelTimer", zap.String(tagTimerID, timerID))
   802  	timerHandle, ok := env.timers[timerID]
   803  	if !ok {
   804  		env.logger.Debug("RequestCancelTimer failed, TimerID not exists.", zap.String(tagTimerID, timerID))
   805  		return
   806  	}
   807  
   808  	delete(env.timers, timerID)
   809  	timerHandle.timer.Stop()
   810  	timerHandle.env.postCallback(func() {
   811  		timerHandle.callback(nil, NewCanceledError())
   812  		if timerHandle.env.onTimerCancelledListener != nil {
   813  			timerHandle.env.onTimerCancelledListener(timerID)
   814  		}
   815  	}, true)
   816  }
   817  
   818  func (env *testWorkflowEnvironmentImpl) Complete(result []byte, err error) {
   819  	if env.isTestCompleted {
   820  		env.logger.Debug("Workflow already completed.")
   821  		return
   822  	}
   823  	env.workflowDef.Close()
   824  	if _, ok := err.(*CanceledError); ok && env.workflowCancelHandler != nil {
   825  		env.workflowCancelHandler()
   826  	}
   827  
   828  	dc := env.GetDataConverter()
   829  	// Test is potentially not over, for parent Cron workflows
   830  	if (!env.isChildWorkflow() && !env.IsCron()) || env.isChildWorkflow() {
   831  		env.isTestCompleted = true
   832  	}
   833  
   834  	if err != nil {
   835  		switch err := err.(type) {
   836  		case *CanceledError, *ContinueAsNewError, *TimeoutError, *shared.WorkflowExecutionAlreadyStartedError:
   837  			env.testError = err
   838  		case *workflowPanicError:
   839  			env.testError = newPanicError(err.value, err.stackTrace)
   840  		default:
   841  			reason, details := getErrorDetails(err, dc)
   842  			env.testError = constructError(reason, details, dc)
   843  		}
   844  	} else {
   845  		env.testResult = newEncodedValue(result, dc)
   846  	}
   847  
   848  	// Only close on:
   849  	// 1. Child-Workflows
   850  	// 2. non-cron Workflows
   851  	if env.isChildWorkflow() && !env.IsCron() {
   852  		close(env.doneChannel)
   853  	}
   854  
   855  	if env.isChildWorkflow() {
   856  		// this is completion of child workflow
   857  		childWorkflowID := env.workflowInfo.WorkflowExecution.ID
   858  		if childWorkflowHandle, ok := env.runningWorkflows[childWorkflowID]; ok && !childWorkflowHandle.handled {
   859  			// It is possible that child workflow could complete after cancellation. In that case, childWorkflowHandle
   860  			// would have already been removed from the runningWorkflows map by RequestCancelWorkflow().
   861  			childWorkflowHandle.handled = true
   862  			// check if a retry is needed
   863  			if childWorkflowHandle.rerun(true) {
   864  				// rerun requested, so we don't want to post the error to parent workflow, return here.
   865  				return
   866  			}
   867  
   868  			// no rerun, child workflow is done.
   869  			env.parentEnv.postCallback(func() {
   870  				// deliver result
   871  				childWorkflowHandle.err = env.testError
   872  				childWorkflowHandle.callback(result, env.testError)
   873  				if env.onChildWorkflowCompletedListener != nil {
   874  					env.onChildWorkflowCompletedListener(env.workflowInfo, env.testResult, env.testError)
   875  				}
   876  			}, true /* true to trigger parent workflow to resume to handle child workflow's result */)
   877  		}
   878  	} else {
   879  		if env.IsCron() {
   880  			workflowID := env.workflowInfo.WorkflowExecution.ID
   881  			if workflowHandle, ok := env.runningWorkflows[workflowID]; ok {
   882  				// On rerun, consider Workflow as not-handled
   883  				if workflowHandle.rerun(false) {
   884  					return
   885  				}
   886  			}
   887  		}
   888  	}
   889  	// No Reruns....Test is Complete
   890  	env.isTestCompleted = true
   891  }
   892  
   893  func (h *testWorkflowHandle) rerun(asChild bool) bool {
   894  	env := h.env
   895  	if asChild && !env.isChildWorkflow() {
   896  		return false
   897  	}
   898  	if !asChild && env.isChildWorkflow() {
   899  		return false
   900  	}
   901  	params := h.params
   902  
   903  	// pass down the last completion result
   904  	var result []byte
   905  	if env.testResult != nil {
   906  		env.testResult.Get(&result)
   907  	}
   908  	if len(result) == 0 {
   909  		// not successful run this time, carry over from whatever previous run pass to this run.
   910  		result = env.workflowInfo.lastCompletionResult
   911  	}
   912  	if asChild {
   913  		params.lastCompletionResult = result
   914  
   915  		if params.retryPolicy != nil && env.testError != nil {
   916  			errReason, _ := getErrorDetails(env.testError, env.GetDataConverter())
   917  			var expireTime time.Time
   918  			if params.retryPolicy.GetExpirationIntervalInSeconds() > 0 {
   919  				expireTime = params.scheduledTime.Add(time.Second * time.Duration(params.retryPolicy.GetExpirationIntervalInSeconds()))
   920  			}
   921  			backoff := getRetryBackoffFromThriftRetryPolicy(params.retryPolicy, env.workflowInfo.Attempt, errReason, env.Now(), expireTime)
   922  			if backoff > 0 {
   923  				// remove the current child workflow from the pending child workflow map because
   924  				// the childWorkflowID will be the same for retry run.
   925  				delete(env.runningWorkflows, env.workflowInfo.WorkflowExecution.ID)
   926  				params.attempt++
   927  				env.parentEnv.executeChildWorkflowWithDelay(backoff, *params, h.callback, nil /* child workflow already started */)
   928  				return true
   929  			}
   930  		}
   931  		if len(params.cronSchedule) > 0 {
   932  			if env.cronMaxIterations < 0 || (env.cronMaxIterations > 0 && env.cronIterations < env.cronMaxIterations) {
   933  				schedule, err := cron.ParseStandard(params.cronSchedule)
   934  				if err != nil {
   935  					panic(fmt.Errorf("invalid cron schedule %v, err: %v", params.cronSchedule, err))
   936  				}
   937  				workflowNow := env.Now().In(time.UTC)
   938  				backoff := schedule.Next(workflowNow).Sub(workflowNow)
   939  				if backoff > 0 {
   940  					env.cronIterations++
   941  					delete(env.runningWorkflows, env.workflowInfo.WorkflowExecution.ID)
   942  					params.attempt = 0
   943  					params.scheduledTime = env.Now()
   944  					env.parentEnv.executeChildWorkflowWithDelay(backoff, *params, h.callback, nil /* child workflow already started */)
   945  					return true
   946  				}
   947  			}
   948  		}
   949  	} else {
   950  		// Re-run a non-Child workflow if it has a Cron Schedule
   951  		if h.env.workflowInfo.CronSchedule != nil {
   952  			if env.cronMaxIterations < 0 || (env.cronMaxIterations > 0 && env.cronIterations < env.cronMaxIterations) {
   953  				cronSchedule := *h.env.workflowInfo.CronSchedule
   954  				if len(cronSchedule) == 0 {
   955  					return false
   956  				}
   957  				schedule, err := cron.ParseStandard(cronSchedule)
   958  				if err != nil {
   959  					panic(fmt.Errorf("invalid cron schedule %v, err: %v", cronSchedule, err))
   960  				}
   961  				workflowNow := env.Now().In(time.UTC)
   962  				backoff := schedule.Next(workflowNow).Sub(workflowNow)
   963  				if backoff > 0 {
   964  					env.cronIterations++
   965  					// Prepare the env for the next iteration
   966  					env.runningCount--
   967  					env.setLastCompletionResult(result)
   968  					// Since MainLoop is already running, we just want to execute the dispatcher
   969  					// which will run the Workflow,
   970  					env.registerDelayedCallback(func() {
   971  						env.runningCount++
   972  						env.workflowDef, _ = env.getWorkflowDefinition(env.workflowInfo.WorkflowType)
   973  						// Use the existing headers and input
   974  						env.workflowDef.Execute(env, env.header, env.workflowInput)
   975  						env.startDecisionTask()
   976  					}, backoff-backoff)
   977  					return true
   978  				}
   979  			}
   980  		}
   981  	}
   982  
   983  	return false
   984  }
   985  
   986  func (env *testWorkflowEnvironmentImpl) CompleteActivity(taskToken []byte, result interface{}, err error) error {
   987  	if taskToken == nil {
   988  		return errors.New("nil task token provided")
   989  	}
   990  	var data []byte
   991  	if result != nil {
   992  		var encodeErr error
   993  		data, encodeErr = encodeArg(env.GetDataConverter(), result)
   994  		if encodeErr != nil {
   995  			return encodeErr
   996  		}
   997  	}
   998  
   999  	activityID := string(taskToken)
  1000  	env.postCallback(func() {
  1001  		activityHandle, ok := env.getActivityHandle(activityID)
  1002  		if !ok {
  1003  			env.logger.Debug("CompleteActivity: ActivityID not found, could be already completed or cancelled.",
  1004  				zap.String(tagActivityID, activityID))
  1005  			return
  1006  		}
  1007  		request := convertActivityResultToRespondRequest("test-identity", taskToken, data, err, env.GetDataConverter())
  1008  		env.handleActivityResult(activityID, request, activityHandle.activityType, env.GetDataConverter())
  1009  	}, false /* do not auto schedule decision task, because activity might be still pending */)
  1010  
  1011  	return nil
  1012  }
  1013  
  1014  func (env *testWorkflowEnvironmentImpl) GetLogger() *zap.Logger {
  1015  	return env.logger
  1016  }
  1017  
  1018  func (env *testWorkflowEnvironmentImpl) GetMetricsScope() tally.Scope {
  1019  	return env.workerOptions.MetricsScope
  1020  }
  1021  
  1022  func (env *testWorkflowEnvironmentImpl) GetDataConverter() DataConverter {
  1023  	return env.workerOptions.DataConverter
  1024  }
  1025  
  1026  func (env *testWorkflowEnvironmentImpl) GetContextPropagators() []ContextPropagator {
  1027  	return env.workerOptions.ContextPropagators
  1028  }
  1029  
  1030  func (env *testWorkflowEnvironmentImpl) ExecuteActivity(parameters executeActivityParams, callback resultHandler) *activityInfo {
  1031  	var activityID string
  1032  	if parameters.ActivityID == nil || *parameters.ActivityID == "" {
  1033  		activityID = getStringID(env.nextID())
  1034  	} else {
  1035  		activityID = *parameters.ActivityID
  1036  	}
  1037  	activityInfo := &activityInfo{activityID: activityID}
  1038  	task := newTestActivityTask(
  1039  		defaultTestWorkflowID,
  1040  		defaultTestRunID,
  1041  		activityInfo.activityID,
  1042  		defaultTestWorkflowTypeName,
  1043  		defaultTestDomainName,
  1044  		parameters,
  1045  	)
  1046  
  1047  	taskHandler := env.newTestActivityTaskHandler(parameters.TaskListName, parameters.DataConverter)
  1048  	activityHandle := &testActivityHandle{callback: callback, activityType: parameters.ActivityType.Name}
  1049  
  1050  	env.setActivityHandle(activityInfo.activityID, activityHandle)
  1051  	env.runningCount++
  1052  	// activity runs in separate goroutinue outside of workflow dispatcher
  1053  	// do callback in a defer to handle calls to runtime.Goexit inside the activity (which is done by t.FailNow)
  1054  	go func() {
  1055  		var result interface{}
  1056  		defer func() {
  1057  			panicErr := recover()
  1058  			if result == nil && panicErr == nil {
  1059  				reason := "activity called runtime.Goexit"
  1060  				result = &shared.RespondActivityTaskFailedRequest{
  1061  					Reason: &reason,
  1062  				}
  1063  			} else if panicErr != nil {
  1064  				reason := errReasonPanic
  1065  				details, _ := env.GetDataConverter().ToData(fmt.Sprintf("%v", panicErr))
  1066  				result = &shared.RespondActivityTaskFailedRequest{
  1067  					Reason:  &reason,
  1068  					Details: details,
  1069  				}
  1070  			}
  1071  			// post activity result to workflow dispatcher
  1072  			env.postCallback(func() {
  1073  				env.handleActivityResult(activityInfo.activityID, result, parameters.ActivityType.Name, parameters.DataConverter)
  1074  				env.runningCount--
  1075  			}, false /* do not auto schedule decision task, because activity might be still pending */)
  1076  		}()
  1077  		result = env.executeActivityWithRetryForTest(taskHandler, parameters, task)
  1078  	}()
  1079  
  1080  	return activityInfo
  1081  }
  1082  
  1083  func (env *testWorkflowEnvironmentImpl) getActivityHandle(activityID string) (*testActivityHandle, bool) {
  1084  	handle, ok := env.activities[env.makeUniqueID(activityID)]
  1085  	return handle, ok
  1086  }
  1087  
  1088  func (env *testWorkflowEnvironmentImpl) setActivityHandle(activityID string, handle *testActivityHandle) {
  1089  	env.activities[env.makeUniqueID(activityID)] = handle
  1090  }
  1091  
  1092  func (env *testWorkflowEnvironmentImpl) deleteHandle(activityID string) {
  1093  	delete(env.activities, env.makeUniqueID(activityID))
  1094  }
  1095  
  1096  func (env *testWorkflowEnvironmentImpl) makeUniqueID(id string) string {
  1097  	// ActivityID is unique per workflow, but different workflow could have same activityID.
  1098  	// Make the key unique globally as we share the same collection for all running workflows in test.
  1099  	return fmt.Sprintf("%v_%v", env.WorkflowInfo().WorkflowExecution.RunID, id)
  1100  }
  1101  
  1102  func (env *testWorkflowEnvironmentImpl) executeActivityWithRetryForTest(
  1103  	taskHandler ActivityTaskHandler,
  1104  	parameters executeActivityParams,
  1105  	task *shared.PollForActivityTaskResponse,
  1106  ) (result interface{}) {
  1107  	var expireTime time.Time
  1108  	if parameters.RetryPolicy != nil && parameters.RetryPolicy.GetExpirationIntervalInSeconds() > 0 {
  1109  		expireTime = env.Now().Add(time.Second * time.Duration(parameters.RetryPolicy.GetExpirationIntervalInSeconds()))
  1110  	}
  1111  
  1112  	for {
  1113  		var err error
  1114  		result, err = taskHandler.Execute(parameters.TaskListName, task)
  1115  		if err != nil {
  1116  			if err == context.DeadlineExceeded {
  1117  				return err
  1118  			}
  1119  			panic(err)
  1120  		}
  1121  
  1122  		// check if a retry is needed
  1123  		if request, ok := result.(*shared.RespondActivityTaskFailedRequest); ok && parameters.RetryPolicy != nil {
  1124  			p := fromThriftRetryPolicy(parameters.RetryPolicy)
  1125  			backoff := getRetryBackoffWithNowTime(p, task.GetAttempt(), *request.Reason, env.Now(), expireTime)
  1126  			if backoff > 0 {
  1127  				// need a retry
  1128  				waitCh := make(chan struct{})
  1129  
  1130  				// register the delayed call back first, otherwise other timers may be fired before the retry timer
  1131  				// is enqueued.
  1132  				env.registerDelayedCallback(func() {
  1133  					env.runningCount++
  1134  					task.Attempt = common.Int32Ptr(task.GetAttempt() + 1)
  1135  					activityID := string(task.TaskToken)
  1136  					if ah, ok := env.getActivityHandle(activityID); ok {
  1137  						task.HeartbeatDetails = ah.heartbeatDetails
  1138  					}
  1139  					close(waitCh)
  1140  				}, backoff)
  1141  				env.postCallback(func() { env.runningCount-- }, false)
  1142  
  1143  				<-waitCh
  1144  				continue
  1145  			}
  1146  		}
  1147  
  1148  		// no retry
  1149  		break
  1150  	}
  1151  
  1152  	return
  1153  }
  1154  
  1155  func fromThriftRetryPolicy(p *shared.RetryPolicy) *RetryPolicy {
  1156  	return &RetryPolicy{
  1157  		InitialInterval:          time.Second * time.Duration(p.GetInitialIntervalInSeconds()),
  1158  		BackoffCoefficient:       p.GetBackoffCoefficient(),
  1159  		MaximumInterval:          time.Second * time.Duration(p.GetMaximumIntervalInSeconds()),
  1160  		ExpirationInterval:       time.Second * time.Duration(p.GetExpirationIntervalInSeconds()),
  1161  		MaximumAttempts:          p.GetMaximumAttempts(),
  1162  		NonRetriableErrorReasons: p.NonRetriableErrorReasons,
  1163  	}
  1164  }
  1165  
  1166  func getRetryBackoffFromThriftRetryPolicy(tp *shared.RetryPolicy, attempt int32, errReason string, now, expireTime time.Time) time.Duration {
  1167  	if tp == nil {
  1168  		return noRetryBackoff
  1169  	}
  1170  
  1171  	p := fromThriftRetryPolicy(tp)
  1172  	return getRetryBackoffWithNowTime(p, attempt, errReason, now, expireTime)
  1173  }
  1174  
  1175  func (env *testWorkflowEnvironmentImpl) ExecuteLocalActivity(params executeLocalActivityParams, callback laResultHandler) *localActivityInfo {
  1176  	activityID := getStringID(env.nextID())
  1177  	wOptions := AugmentWorkerOptions(env.workerOptions)
  1178  	ae := &activityExecutor{name: getActivityFunctionName(env.registry, params.ActivityFn), fn: params.ActivityFn}
  1179  	if at, _ := getValidatedActivityFunction(params.ActivityFn, params.InputArgs, env.registry); at != nil {
  1180  		// local activity could be registered, if so use the registered name. This name is only used to find a mock.
  1181  		ae.name = at.Name
  1182  	}
  1183  	aew := &activityExecutorWrapper{activityExecutor: ae, env: env}
  1184  
  1185  	// substitute the local activity function so we could replace with mock if it is supplied.
  1186  	params.ActivityFn = func(ctx context.Context, inputArgs ...interface{}) ([]byte, error) {
  1187  		return aew.ExecuteWithActualArgs(ctx, params.InputArgs)
  1188  	}
  1189  
  1190  	task := newLocalActivityTask(params, callback, activityID)
  1191  	taskHandler := localActivityTaskHandler{
  1192  		userContext:        wOptions.BackgroundActivityContext,
  1193  		metricsScope:       metrics.NewTaggedScope(wOptions.MetricsScope),
  1194  		logger:             wOptions.Logger,
  1195  		dataConverter:      wOptions.DataConverter,
  1196  		tracer:             wOptions.Tracer,
  1197  		contextPropagators: wOptions.ContextPropagators,
  1198  	}
  1199  
  1200  	env.localActivities[activityID] = task
  1201  	env.runningCount++
  1202  
  1203  	go func() {
  1204  		result := taskHandler.executeLocalActivityTask(task)
  1205  		env.postCallback(func() {
  1206  			env.handleLocalActivityResult(result)
  1207  			env.runningCount--
  1208  		}, false)
  1209  	}()
  1210  
  1211  	return &localActivityInfo{activityID: activityID}
  1212  }
  1213  
  1214  func (env *testWorkflowEnvironmentImpl) RequestCancelLocalActivity(activityID string) {
  1215  	task, ok := env.localActivities[activityID]
  1216  	if !ok {
  1217  		env.logger.Debug("RequestCancelLocalActivity failed, LocalActivity not exists or already completed.", zap.String(tagActivityID, activityID))
  1218  		return
  1219  	}
  1220  	activityInfo := env.getActivityInfo(activityID, getActivityFunctionName(env.registry, task.params.ActivityFn))
  1221  	env.logger.Debug("RequestCancelLocalActivity", zap.String(tagActivityID, activityID))
  1222  	delete(env.localActivities, activityID)
  1223  	env.postCallback(func() {
  1224  		lar := &localActivityResultWrapper{err: ErrCanceled, backoff: noRetryBackoff}
  1225  		task.callback(lar)
  1226  		if env.onLocalActivityCanceledListener != nil {
  1227  			env.onLocalActivityCanceledListener(activityInfo)
  1228  		}
  1229  	}, true)
  1230  }
  1231  
  1232  func (env *testWorkflowEnvironmentImpl) handleActivityResult(activityID string, result interface{}, activityType string,
  1233  	dataConverter DataConverter) {
  1234  	env.logger.Debug(fmt.Sprintf("handleActivityResult: %T.", result),
  1235  		zap.String(tagActivityID, activityID), zap.String(tagActivityType, activityType))
  1236  	activityInfo := env.getActivityInfo(activityID, activityType)
  1237  	if result == ErrActivityResultPending {
  1238  		// In case activity returns ErrActivityResultPending, the respond will be nil, and we don't need to do anything.
  1239  		// Activity will need to complete asynchronously using CompleteActivity().
  1240  		if env.onActivityCompletedListener != nil {
  1241  			env.onActivityCompletedListener(activityInfo, nil, ErrActivityResultPending)
  1242  		}
  1243  		return
  1244  	}
  1245  
  1246  	// this is running in dispatcher
  1247  	activityHandle, ok := env.getActivityHandle(activityID)
  1248  	if !ok {
  1249  		env.logger.Debug("handleActivityResult: ActivityID not exists, could be already completed or cancelled.",
  1250  			zap.String(tagActivityID, activityID))
  1251  		return
  1252  	}
  1253  
  1254  	delete(env.activities, activityID)
  1255  
  1256  	var blob []byte
  1257  	var err error
  1258  
  1259  	switch request := result.(type) {
  1260  	case *shared.RespondActivityTaskCanceledRequest:
  1261  		details := newEncodedValues(request.Details, dataConverter)
  1262  		err = NewCanceledError(details)
  1263  		activityHandle.callback(nil, err)
  1264  	case *shared.RespondActivityTaskFailedRequest:
  1265  		err = constructError(*request.Reason, request.Details, dataConverter)
  1266  		activityHandle.callback(nil, err)
  1267  	case *shared.RespondActivityTaskCompletedRequest:
  1268  		blob = request.Result
  1269  		activityHandle.callback(blob, nil)
  1270  	default:
  1271  		if result == context.DeadlineExceeded {
  1272  			err = NewTimeoutError(shared.TimeoutTypeStartToClose, context.DeadlineExceeded.Error())
  1273  			activityHandle.callback(nil, err)
  1274  		} else {
  1275  			panic(fmt.Sprintf("unsupported respond type %T", result))
  1276  		}
  1277  	}
  1278  
  1279  	if env.onActivityCompletedListener != nil {
  1280  		if err != nil {
  1281  			env.onActivityCompletedListener(activityInfo, nil, err)
  1282  		} else {
  1283  			env.onActivityCompletedListener(activityInfo, newEncodedValue(blob, dataConverter), nil)
  1284  		}
  1285  	}
  1286  
  1287  	env.startDecisionTask()
  1288  }
  1289  
  1290  func (env *testWorkflowEnvironmentImpl) handleLocalActivityResult(result *localActivityResult) {
  1291  	activityID := result.task.activityID
  1292  	activityType := getActivityFunctionName(env.registry, result.task.params.ActivityFn)
  1293  	env.logger.Debug(fmt.Sprintf("handleLocalActivityResult: Err: %v, Result: %v.", result.err, string(result.result)),
  1294  		zap.String(tagActivityID, activityID), zap.String(tagActivityType, activityType))
  1295  
  1296  	activityInfo := env.getActivityInfo(activityID, activityType)
  1297  	task, ok := env.localActivities[activityID]
  1298  	if !ok {
  1299  		env.logger.Debug("handleLocalActivityResult: ActivityID not exists, could be already completed or cancelled.",
  1300  			zap.String(tagActivityID, activityID))
  1301  		return
  1302  	}
  1303  
  1304  	delete(env.localActivities, activityID)
  1305  	var encodedErr error
  1306  	if result.err != nil {
  1307  		errReason, errDetails := getErrorDetails(result.err, env.GetDataConverter())
  1308  		encodedErr = constructError(errReason, errDetails, env.GetDataConverter())
  1309  	}
  1310  	lar := &localActivityResultWrapper{
  1311  		err:     encodedErr,
  1312  		result:  result.result,
  1313  		backoff: noRetryBackoff,
  1314  	}
  1315  	if result.task.retryPolicy != nil && result.err != nil {
  1316  		lar.backoff = getRetryBackoff(result, env.Now())
  1317  		lar.attempt = task.attempt
  1318  	}
  1319  	task.callback(lar)
  1320  	if env.onLocalActivityCompletedListener != nil {
  1321  		if result.err != nil {
  1322  			env.onLocalActivityCompletedListener(activityInfo, nil, result.err)
  1323  		} else {
  1324  			env.onLocalActivityCompletedListener(activityInfo, newEncodedValue(result.result, env.GetDataConverter()), nil)
  1325  		}
  1326  	}
  1327  
  1328  	env.startDecisionTask()
  1329  }
  1330  
  1331  // runBeforeMockCallReturns is registered as mock call's RunFn by *mock.Call.Run(fn). It will be called by testify's
  1332  // mock.MethodCalled() before it returns.
  1333  func (env *testWorkflowEnvironmentImpl) runBeforeMockCallReturns(call *MockCallWrapper, args mock.Arguments) {
  1334  	var waitDuration time.Duration
  1335  	if call.waitDuration != nil {
  1336  		waitDuration = call.waitDuration()
  1337  	}
  1338  	if waitDuration > 0 {
  1339  		// we want this mock call to block until the wait duration is elapsed (on workflow clock).
  1340  		waitCh := make(chan time.Time)
  1341  		env.registerDelayedCallback(func() {
  1342  			env.runningCount++  // increase runningCount as the mock call is ready to resume.
  1343  			waitCh <- env.Now() // this will unblock mock call
  1344  		}, waitDuration)
  1345  
  1346  		// make sure decrease runningCount after delayed callback is posted
  1347  		env.postCallback(func() {
  1348  			env.runningCount-- // reduce runningCount, since this mock call is about to be blocked.
  1349  		}, false)
  1350  		<-waitCh // this will block until mock clock move forward by waitDuration
  1351  	}
  1352  
  1353  	// run the actual runFn if it was setup
  1354  	if call.runFn != nil {
  1355  		call.runFn(args)
  1356  	}
  1357  }
  1358  
  1359  // Execute executes the activity code.
  1360  func (a *activityExecutorWrapper) Execute(ctx context.Context, input []byte) ([]byte, error) {
  1361  	activityInfo := GetActivityInfo(ctx)
  1362  	dc := getDataConverterFromActivityCtx(ctx)
  1363  	if a.env.onActivityStartedListener != nil {
  1364  		waitCh := make(chan struct{})
  1365  		a.env.postCallback(func() {
  1366  			a.env.onActivityStartedListener(&activityInfo, ctx, newEncodedValues(input, dc))
  1367  			close(waitCh)
  1368  		}, false)
  1369  		<-waitCh // wait until listener returns
  1370  	}
  1371  
  1372  	m := &mockWrapper{env: a.env, name: a.name, fn: a.fn, isWorkflow: false, dataConverter: dc}
  1373  	if mockRet := m.getMockReturn(ctx, input); mockRet != nil {
  1374  		return m.executeMock(ctx, input, mockRet)
  1375  	}
  1376  
  1377  	return a.activityExecutor.Execute(ctx, input)
  1378  }
  1379  
  1380  // Execute executes the activity code.
  1381  func (a *activityExecutorWrapper) ExecuteWithActualArgs(ctx context.Context, inputArgs []interface{}) ([]byte, error) {
  1382  	activityInfo := GetActivityInfo(ctx)
  1383  	if a.env.onLocalActivityStartedListener != nil {
  1384  		waitCh := make(chan struct{})
  1385  		a.env.postCallback(func() {
  1386  			a.env.onLocalActivityStartedListener(&activityInfo, ctx, inputArgs)
  1387  			close(waitCh)
  1388  		}, false)
  1389  		<-waitCh
  1390  	}
  1391  
  1392  	m := &mockWrapper{env: a.env, name: a.name, fn: a.fn, isWorkflow: false}
  1393  	if mockRet := m.getMockReturnWithActualArgs(ctx, inputArgs); mockRet != nil {
  1394  		return m.executeMockWithActualArgs(ctx, inputArgs, mockRet)
  1395  	}
  1396  
  1397  	return a.activityExecutor.ExecuteWithActualArgs(ctx, inputArgs)
  1398  }
  1399  
  1400  // Execute executes the workflow code.
  1401  func (w *workflowExecutorWrapper) Execute(ctx Context, input []byte) (result []byte, err error) {
  1402  	env := w.env
  1403  	if env.isChildWorkflow() && env.onChildWorkflowStartedListener != nil {
  1404  		env.onChildWorkflowStartedListener(GetWorkflowInfo(ctx), ctx, newEncodedValues(input, w.env.GetDataConverter()))
  1405  	}
  1406  
  1407  	if !env.isChildWorkflow() {
  1408  		// This is to prevent auto-forwarding mock clock before main workflow starts. For child workflow, we increase
  1409  		// the counter in env.ExecuteChildWorkflow(). We cannot do it here for child workflow, because we need to make
  1410  		// sure the counter is increased before returning from ExecuteChildWorkflow().
  1411  		env.runningCount++
  1412  	}
  1413  
  1414  	m := &mockWrapper{
  1415  		env:           env,
  1416  		name:          w.workflowType,
  1417  		fn:            w.fn,
  1418  		isWorkflow:    true,
  1419  		dataConverter: env.GetDataConverter(),
  1420  	}
  1421  	// This method is called by workflow's dispatcher. In this test suite, it is run in the main loop. We cannot block
  1422  	// the main loop, but the mock could block if it is configured to wait. So we need to use a separate goroutinue to
  1423  	// run the mock, and resume after mock call returns.
  1424  	mockReadyChannel := NewChannel(ctx)
  1425  	// Make a copy of the context for getMockReturn() call to avoid race condition.
  1426  	// Use existing interceptors from env.
  1427  	envInterceptor := &workflowEnvironmentInterceptor{env: env}
  1428  	ctxCopy := newWorkflowContext(w.env, envInterceptor, envInterceptor)
  1429  	go func() {
  1430  		// getMockReturn could block if mock is configured to wait. The returned mockRet is what has been configured
  1431  		// for the mock by using MockCallWrapper.Return(). The mockRet could be mock values or mock function. We process
  1432  		// the returned mockRet by calling executeMock() later in the main thread after it is send over via mockReadyChannel.
  1433  		mockRet := m.getMockReturn(ctxCopy, input)
  1434  		env.postCallback(func() {
  1435  			mockReadyChannel.SendAsync(mockRet)
  1436  		}, true /* true to trigger the dispatcher for this workflow so it resume from mockReadyChannel block*/)
  1437  	}()
  1438  
  1439  	var mockRet mock.Arguments
  1440  	// This will block workflow dispatcher (on cadence channel), which the dispatcher understand and will return from
  1441  	// ExecuteUntilAllBlocked() so the main loop is not blocked. The dispatcher will unblock when getMockReturn() returns.
  1442  	mockReadyChannel.Receive(ctx, &mockRet)
  1443  
  1444  	// reduce runningCount to allow auto-forwarding mock clock after current workflow dispatcher run is blocked (aka
  1445  	// ExecuteUntilAllBlocked() returns).
  1446  	env.runningCount--
  1447  
  1448  	childWE := env.workflowInfo.WorkflowExecution
  1449  	var startedErr error
  1450  	if mockRet != nil {
  1451  		// workflow was mocked.
  1452  		result, err = m.executeMock(ctx, input, mockRet)
  1453  		if env.isChildWorkflow() && err == ErrMockStartChildWorkflowFailed {
  1454  			childWE, startedErr = WorkflowExecution{}, err
  1455  		}
  1456  	}
  1457  
  1458  	if env.isChildWorkflow() && env.startedHandler != nil /* startedHandler could be nil for retry */ {
  1459  		// notify parent that child workflow is started
  1460  		env.parentEnv.postCallback(func() {
  1461  			env.startedHandler(childWE, startedErr)
  1462  		}, true)
  1463  	}
  1464  
  1465  	if mockRet != nil {
  1466  		return result, err
  1467  	}
  1468  
  1469  	// no mock, so call the actual workflow
  1470  	return w.workflowExecutor.Execute(ctx, input)
  1471  }
  1472  
  1473  func (m *mockWrapper) getCtxArg(ctx interface{}) []interface{} {
  1474  	fnType := reflect.TypeOf(m.fn)
  1475  	if fnType.NumIn() > 0 {
  1476  		if (!m.isWorkflow && isActivityContext(fnType.In(0))) ||
  1477  			(m.isWorkflow && isWorkflowContext(fnType.In(0))) {
  1478  			return []interface{}{ctx}
  1479  		}
  1480  	}
  1481  	return nil
  1482  }
  1483  
  1484  func (m *mockWrapper) getMockReturn(ctx interface{}, input []byte) (retArgs mock.Arguments) {
  1485  	if _, ok := m.env.expectedMockCalls[m.name]; !ok {
  1486  		// no mock
  1487  		return nil
  1488  	}
  1489  
  1490  	fnType := reflect.TypeOf(m.fn)
  1491  	reflectArgs, err := decodeArgs(m.dataConverter, fnType, input)
  1492  	if err != nil {
  1493  		panic(fmt.Sprintf("Decode error: %v in %v of type %T", err.Error(), m.name, m.fn))
  1494  	}
  1495  	realArgs := m.getCtxArg(ctx)
  1496  	for _, arg := range reflectArgs {
  1497  		realArgs = append(realArgs, arg.Interface())
  1498  	}
  1499  
  1500  	return m.env.mock.MethodCalled(m.name, realArgs...)
  1501  }
  1502  
  1503  func (m *mockWrapper) getMockReturnWithActualArgs(ctx interface{}, inputArgs []interface{}) (retArgs mock.Arguments) {
  1504  	if _, ok := m.env.expectedMockCalls[m.name]; !ok {
  1505  		// no mock
  1506  		return nil
  1507  	}
  1508  
  1509  	realArgs := m.getCtxArg(ctx)
  1510  	realArgs = append(realArgs, inputArgs...)
  1511  	return m.env.mock.MethodCalled(m.name, realArgs...)
  1512  }
  1513  
  1514  func (m *mockWrapper) getMockFn(mockRet mock.Arguments) interface{} {
  1515  	fnName := m.name
  1516  	mockRetLen := len(mockRet)
  1517  	if mockRetLen == 0 {
  1518  		panic(fmt.Sprintf("mock of %v has no returns", fnName))
  1519  	}
  1520  
  1521  	fnType := reflect.TypeOf(m.fn)
  1522  	// check if mock returns function which must match to the actual function.
  1523  	mockFn := mockRet.Get(0)
  1524  	mockFnType := reflect.TypeOf(mockFn)
  1525  	if mockFnType != nil && mockFnType.Kind() == reflect.Func {
  1526  		if mockFnType != fnType {
  1527  			panic(fmt.Sprintf("mock of %v has incorrect return function, expected %v, but actual is %v",
  1528  				fnName, fnType, mockFnType))
  1529  		}
  1530  		return mockFn
  1531  	}
  1532  	return nil
  1533  }
  1534  
  1535  func (m *mockWrapper) getMockValue(mockRet mock.Arguments) ([]byte, error) {
  1536  	fnName := m.name
  1537  	mockRetLen := len(mockRet)
  1538  	fnType := reflect.TypeOf(m.fn)
  1539  	// check if mockRet have same types as function's return types
  1540  	if mockRetLen != fnType.NumOut() {
  1541  		panic(fmt.Sprintf("mock of %v has incorrect number of returns, expected %d, but actual is %d",
  1542  			fnName, fnType.NumOut(), mockRetLen))
  1543  	}
  1544  	// we already verified function either has 1 return value (error) or 2 return values (result, error)
  1545  	var retErr error
  1546  	mockErr := mockRet[mockRetLen-1] // last mock return must be error
  1547  	if mockErr == nil {
  1548  		retErr = nil
  1549  	} else if err, ok := mockErr.(error); ok {
  1550  		retErr = err
  1551  	} else {
  1552  		panic(fmt.Sprintf("mock of %v has incorrect return type, expected error, but actual is %T (%v)",
  1553  			fnName, mockErr, mockErr))
  1554  	}
  1555  
  1556  	switch mockRetLen {
  1557  	case 1:
  1558  		return nil, retErr
  1559  	case 2:
  1560  		expectedType := fnType.Out(0)
  1561  		mockResult := mockRet[0]
  1562  		if mockResult == nil {
  1563  			switch expectedType.Kind() {
  1564  			case reflect.Ptr, reflect.Interface, reflect.Map, reflect.Slice, reflect.Array:
  1565  				// these are supported nil-able types. (reflect.Chan, reflect.Func are nil-able, but not supported)
  1566  				return nil, retErr
  1567  			default:
  1568  				panic(fmt.Sprintf("mock of %v has incorrect return type, expected %v, but actual is %T (%v)",
  1569  					fnName, expectedType, mockResult, mockResult))
  1570  			}
  1571  		} else {
  1572  			if !reflect.TypeOf(mockResult).AssignableTo(expectedType) {
  1573  				panic(fmt.Sprintf("mock of %v has incorrect return type, expected %v, but actual is %T (%v)",
  1574  					fnName, expectedType, mockResult, mockResult))
  1575  			}
  1576  			result, encodeErr := encodeArg(m.env.GetDataConverter(), mockResult)
  1577  			if encodeErr != nil {
  1578  				panic(fmt.Sprintf("encode result from mock of %v failed: %v", fnName, encodeErr))
  1579  			}
  1580  			return result, retErr
  1581  		}
  1582  	default:
  1583  		// this will never happen, panic just in case
  1584  		panic("mock should either have 1 return value (error) or 2 return values (result, error)")
  1585  	}
  1586  }
  1587  
  1588  func (m *mockWrapper) executeMock(ctx interface{}, input []byte, mockRet mock.Arguments) (result []byte, err error) {
  1589  	// have to handle panics here to support calling ExecuteChildWorkflow(...).GetChildWorkflowExecution().Get(...)
  1590  	// when a child is mocked.
  1591  	defer func() {
  1592  		if r := recover(); r != nil {
  1593  			st := getStackTrace("executeMock", "panic", 4)
  1594  			err = newPanicError(r, st)
  1595  		}
  1596  	}()
  1597  
  1598  	fnName := m.name
  1599  	// check if mock returns function which must match to the actual function.
  1600  	if mockFn := m.getMockFn(mockRet); mockFn != nil {
  1601  		// we found a mock function that matches to actual function, so call that mockFn
  1602  		if m.isWorkflow {
  1603  			executor := &workflowExecutor{workflowType: fnName, fn: mockFn}
  1604  			return executor.Execute(ctx.(Context), input)
  1605  		}
  1606  		executor := &activityExecutor{name: fnName, fn: mockFn}
  1607  		return executor.Execute(ctx.(context.Context), input)
  1608  	}
  1609  
  1610  	return m.getMockValue(mockRet)
  1611  }
  1612  
  1613  func (m *mockWrapper) executeMockWithActualArgs(ctx interface{}, inputArgs []interface{}, mockRet mock.Arguments) ([]byte, error) {
  1614  	fnName := m.name
  1615  	// check if mock returns function which must match to the actual function.
  1616  	if mockFn := m.getMockFn(mockRet); mockFn != nil {
  1617  		executor := &activityExecutor{name: fnName, fn: mockFn}
  1618  		return executor.ExecuteWithActualArgs(ctx.(context.Context), inputArgs)
  1619  	}
  1620  
  1621  	return m.getMockValue(mockRet)
  1622  }
  1623  
  1624  func (env *testWorkflowEnvironmentImpl) newTestActivityTaskHandler(taskList string, dataConverter DataConverter) ActivityTaskHandler {
  1625  	wOptions := AugmentWorkerOptions(env.workerOptions)
  1626  	wOptions.DataConverter = dataConverter
  1627  	params := workerExecutionParameters{
  1628  		WorkerOptions:     wOptions,
  1629  		TaskList:          taskList,
  1630  		UserContext:       wOptions.BackgroundActivityContext,
  1631  		WorkerStopChannel: env.workerStopChannel,
  1632  	}
  1633  	ensureRequiredParams(&params)
  1634  	if params.UserContext == nil {
  1635  		params.UserContext = context.Background()
  1636  	}
  1637  	if env.sessionEnvironment == nil {
  1638  		env.sessionEnvironment = newTestSessionEnvironment(env, &params, wOptions.MaxConcurrentSessionExecutionSize)
  1639  	}
  1640  	params.UserContext = context.WithValue(params.UserContext, sessionEnvironmentContextKey, env.sessionEnvironment)
  1641  	registry := env.registry
  1642  	if len(registry.getRegisteredActivities()) == 0 {
  1643  		panic(fmt.Sprintf("no activity is registered for tasklist '%v'", taskList))
  1644  	}
  1645  
  1646  	getActivity := func(name string) activity {
  1647  		tlsa, ok := env.taskListSpecificActivities[name]
  1648  		if ok {
  1649  			_, ok := tlsa.taskLists[taskList]
  1650  			if !ok {
  1651  				// activity are bind to specific task list but not to current task list
  1652  				return nil
  1653  			}
  1654  		}
  1655  
  1656  		activity, ok := registry.GetActivity(name)
  1657  		if !ok {
  1658  			return nil
  1659  		}
  1660  		ae := &activityExecutor{name: activity.ActivityType().Name, fn: activity.GetFunction()}
  1661  
  1662  		// Special handling for session creation and completion activities.
  1663  		// If real creation activity is used, it will block timers from autofiring.
  1664  		if ae.name == sessionCreationActivityName {
  1665  			ae.fn = sessionCreationActivityForTest
  1666  		}
  1667  		if ae.name == sessionCompletionActivityName {
  1668  			ae.fn = sessionCompletionActivityForTest
  1669  		}
  1670  
  1671  		return &activityExecutorWrapper{activityExecutor: ae, env: env}
  1672  	}
  1673  
  1674  	taskHandler := newActivityTaskHandlerWithCustomProvider(env.service, params, registry, getActivity)
  1675  	return taskHandler
  1676  }
  1677  
  1678  func newTestActivityTask(workflowID, runID, activityID, workflowTypeName, domainName string, params executeActivityParams) *shared.PollForActivityTaskResponse {
  1679  	task := &shared.PollForActivityTaskResponse{
  1680  		WorkflowExecution: &shared.WorkflowExecution{
  1681  			WorkflowId: common.StringPtr(workflowID),
  1682  			RunId:      common.StringPtr(runID),
  1683  		},
  1684  		ActivityId:                      common.StringPtr(activityID),
  1685  		TaskToken:                       []byte(activityID), // use activityID as TaskToken so we can map TaskToken in heartbeat calls.
  1686  		ActivityType:                    &shared.ActivityType{Name: common.StringPtr(params.ActivityType.Name)},
  1687  		Input:                           params.Input,
  1688  		ScheduledTimestamp:              common.Int64Ptr(time.Now().UnixNano()),
  1689  		ScheduleToCloseTimeoutSeconds:   common.Int32Ptr(params.ScheduleToCloseTimeoutSeconds),
  1690  		ScheduledTimestampOfThisAttempt: common.Int64Ptr(time.Now().UnixNano()),
  1691  		StartedTimestamp:                common.Int64Ptr(time.Now().UnixNano()),
  1692  		StartToCloseTimeoutSeconds:      common.Int32Ptr(params.StartToCloseTimeoutSeconds),
  1693  		HeartbeatTimeoutSeconds:         common.Int32Ptr(params.HeartbeatTimeoutSeconds),
  1694  		WorkflowType: &shared.WorkflowType{
  1695  			Name: common.StringPtr(workflowTypeName),
  1696  		},
  1697  		WorkflowDomain: common.StringPtr(domainName),
  1698  		Header:         params.Header,
  1699  	}
  1700  	return task
  1701  }
  1702  
  1703  func (env *testWorkflowEnvironmentImpl) newTimer(d time.Duration, callback resultHandler, notifyListener bool) *timerInfo {
  1704  	nextID := env.nextID()
  1705  	timerInfo := &timerInfo{timerID: getStringID(nextID)}
  1706  	timer := env.mockClock.AfterFunc(d, func() {
  1707  		delete(env.timers, timerInfo.timerID)
  1708  		env.postCallback(func() {
  1709  			callback(nil, nil)
  1710  			if notifyListener && env.onTimerFiredListener != nil {
  1711  				env.onTimerFiredListener(timerInfo.timerID)
  1712  			}
  1713  		}, true)
  1714  	})
  1715  	env.timers[timerInfo.timerID] = &testTimerHandle{
  1716  		env:            env,
  1717  		callback:       callback,
  1718  		timer:          timer,
  1719  		mockTimeToFire: env.mockClock.Now().Add(d),
  1720  		wallTimeToFire: env.wallClock.Now().Add(d),
  1721  		duration:       d,
  1722  		timerID:        nextID,
  1723  	}
  1724  	if notifyListener && env.onTimerScheduledListener != nil {
  1725  		env.onTimerScheduledListener(timerInfo.timerID, d)
  1726  	}
  1727  	return timerInfo
  1728  }
  1729  
  1730  func (env *testWorkflowEnvironmentImpl) NewTimer(d time.Duration, callback resultHandler) *timerInfo {
  1731  	return env.newTimer(d, callback, true)
  1732  }
  1733  
  1734  func (env *testWorkflowEnvironmentImpl) Now() time.Time {
  1735  	return env.mockClock.Now()
  1736  }
  1737  
  1738  func (env *testWorkflowEnvironmentImpl) WorkflowInfo() *WorkflowInfo {
  1739  	return env.workflowInfo
  1740  }
  1741  
  1742  func (env *testWorkflowEnvironmentImpl) RegisterWorkflow(w interface{}) {
  1743  	env.registry.RegisterWorkflow(w)
  1744  }
  1745  
  1746  func (env *testWorkflowEnvironmentImpl) RegisterWorkflowWithOptions(w interface{}, options RegisterWorkflowOptions) {
  1747  	env.registry.RegisterWorkflowWithOptions(w, options)
  1748  }
  1749  
  1750  func (env *testWorkflowEnvironmentImpl) RegisterActivity(a interface{}) {
  1751  	env.registry.RegisterActivity(a)
  1752  }
  1753  
  1754  func (env *testWorkflowEnvironmentImpl) RegisterActivityWithOptions(a interface{}, options RegisterActivityOptions) {
  1755  	env.registry.RegisterActivityWithOptions(a, options)
  1756  }
  1757  
  1758  func (env *testWorkflowEnvironmentImpl) RegisterCancelHandler(handler func()) {
  1759  	env.workflowCancelHandler = handler
  1760  }
  1761  
  1762  func (env *testWorkflowEnvironmentImpl) RegisterSignalHandler(handler func(name string, input []byte)) {
  1763  	env.signalHandler = handler
  1764  }
  1765  
  1766  func (env *testWorkflowEnvironmentImpl) RegisterQueryHandler(handler func(string, []byte) ([]byte, error)) {
  1767  	env.queryHandler = handler
  1768  }
  1769  
  1770  func (env *testWorkflowEnvironmentImpl) RequestCancelChildWorkflow(domainName, workflowID string) {
  1771  	if childHandle, ok := env.runningWorkflows[workflowID]; ok && !childHandle.handled {
  1772  		// current workflow is a parent workflow, and we are canceling a child workflow
  1773  		childEnv := childHandle.env
  1774  		childEnv.cancelWorkflow(func(result []byte, err error) {})
  1775  		return
  1776  	}
  1777  }
  1778  
  1779  func (env *testWorkflowEnvironmentImpl) RequestCancelExternalWorkflow(domainName, workflowID, runID string, callback resultHandler) {
  1780  	if env.workflowInfo.WorkflowExecution.ID == workflowID {
  1781  		// cancel current workflow
  1782  		env.workflowCancelHandler()
  1783  		// check if current workflow is a child workflow
  1784  		if env.isChildWorkflow() && env.onChildWorkflowCanceledListener != nil {
  1785  			env.postCallback(func() {
  1786  				env.onChildWorkflowCanceledListener(env.workflowInfo)
  1787  			}, false)
  1788  		}
  1789  		return
  1790  	} else if childHandle, ok := env.runningWorkflows[workflowID]; ok && !childHandle.handled {
  1791  		// current workflow is a parent workflow, and we are canceling a child workflow
  1792  		if !childHandle.params.waitForCancellation {
  1793  			childHandle.env.Complete(nil, ErrCanceled)
  1794  		}
  1795  		childEnv := childHandle.env
  1796  		env.postCallback(func() {
  1797  			callback(nil, nil)
  1798  		}, true)
  1799  		childEnv.cancelWorkflow(callback)
  1800  		return
  1801  	}
  1802  
  1803  	// target workflow is not child workflow, we need the mock. The mock needs to be called in a separate goroutinue
  1804  	// so it can block and wait on the requested delay time (if configured). If we run it in main thread, and the mock
  1805  	// configured to delay, it will block the main loop which stops the world.
  1806  	env.runningCount++
  1807  	go func() {
  1808  		args := []interface{}{domainName, workflowID, runID}
  1809  		// below call will panic if mock is not properly setup.
  1810  		mockRet := env.mock.MethodCalled(mockMethodForRequestCancelExternalWorkflow, args...)
  1811  		m := &mockWrapper{name: mockMethodForRequestCancelExternalWorkflow, fn: mockFnRequestCancelExternalWorkflow}
  1812  		var err error
  1813  		if mockFn := m.getMockFn(mockRet); mockFn != nil {
  1814  			executor := &activityExecutor{name: mockMethodForRequestCancelExternalWorkflow, fn: mockFn}
  1815  			_, err = executor.ExecuteWithActualArgs(nil, args)
  1816  		} else {
  1817  			_, err = m.getMockValue(mockRet)
  1818  		}
  1819  		env.postCallback(func() {
  1820  			callback(nil, err)
  1821  			env.runningCount--
  1822  		}, true)
  1823  	}()
  1824  }
  1825  
  1826  func (env *testWorkflowEnvironmentImpl) IsReplaying() bool {
  1827  	// this test environment never replay
  1828  	return false
  1829  }
  1830  
  1831  func (env *testWorkflowEnvironmentImpl) IsCron() bool {
  1832  	// this test environment never replay
  1833  	return env.workflowInfo.CronSchedule != nil && len(*env.workflowInfo.CronSchedule) > 0
  1834  }
  1835  
  1836  func (env *testWorkflowEnvironmentImpl) SignalExternalWorkflow(domainName, workflowID, runID, signalName string, input []byte, arg interface{}, childWorkflowOnly bool, callback resultHandler) {
  1837  	// check if target workflow is a known workflow
  1838  	if childHandle, ok := env.runningWorkflows[workflowID]; ok {
  1839  		// target workflow is a child
  1840  		childEnv := childHandle.env
  1841  		if childEnv.isTestCompleted {
  1842  			// child already completed (NOTE: we have only one failed cause now)
  1843  			err := newUnknownExternalWorkflowExecutionError()
  1844  			callback(nil, err)
  1845  		} else {
  1846  			childEnv.signalHandler(signalName, input)
  1847  			callback(nil, nil)
  1848  		}
  1849  		childEnv.postCallback(func() {}, true) // resume child workflow since a signal is sent.
  1850  		return
  1851  	}
  1852  
  1853  	// here we signal a child workflow but we cannot find it
  1854  	if childWorkflowOnly {
  1855  		err := newUnknownExternalWorkflowExecutionError()
  1856  		callback(nil, err)
  1857  		return
  1858  	}
  1859  
  1860  	// target workflow is not child workflow, we need the mock. The mock needs to be called in a separate goroutinue
  1861  	// so it can block and wait on the requested delay time (if configured). If we run it in main thread, and the mock
  1862  	// configured to delay, it will block the main loop which stops the world.
  1863  	env.runningCount++
  1864  	go func() {
  1865  		args := []interface{}{domainName, workflowID, runID, signalName, arg}
  1866  		// below call will panic if mock is not properly setup.
  1867  		mockRet := env.mock.MethodCalled(mockMethodForSignalExternalWorkflow, args...)
  1868  		m := &mockWrapper{name: mockMethodForSignalExternalWorkflow, fn: mockFnSignalExternalWorkflow}
  1869  		var err error
  1870  		if mockFn := m.getMockFn(mockRet); mockFn != nil {
  1871  			executor := &activityExecutor{name: mockMethodForSignalExternalWorkflow, fn: mockFn}
  1872  			_, err = executor.ExecuteWithActualArgs(nil, args)
  1873  		} else {
  1874  			_, err = m.getMockValue(mockRet)
  1875  		}
  1876  		env.postCallback(func() {
  1877  			callback(nil, err)
  1878  			env.runningCount--
  1879  		}, true)
  1880  	}()
  1881  }
  1882  
  1883  func (env *testWorkflowEnvironmentImpl) ExecuteChildWorkflow(params executeWorkflowParams, callback resultHandler, startedHandler func(r WorkflowExecution, e error)) error {
  1884  	return env.executeChildWorkflowWithDelay(0, params, callback, startedHandler)
  1885  }
  1886  
  1887  func (env *testWorkflowEnvironmentImpl) executeChildWorkflowWithDelay(delayStart time.Duration, params executeWorkflowParams, callback resultHandler, startedHandler func(r WorkflowExecution, e error)) error {
  1888  	childEnv, err := env.newTestWorkflowEnvironmentForChild(&params, callback, startedHandler)
  1889  	if err != nil {
  1890  		env.logger.Sugar().Infof("ExecuteChildWorkflow failed: %v", err)
  1891  		return err
  1892  	}
  1893  
  1894  	env.logger.Sugar().Infof("ExecuteChildWorkflow: %v", params.workflowType.Name)
  1895  	env.runningCount++
  1896  
  1897  	// run child workflow in separate goroutinue
  1898  	go childEnv.executeWorkflowInternal(delayStart, params.workflowType.Name, params.input)
  1899  
  1900  	return nil
  1901  }
  1902  
  1903  func (env *testWorkflowEnvironmentImpl) SideEffect(f func() ([]byte, error), callback resultHandler) {
  1904  	callback(f())
  1905  }
  1906  
  1907  func (env *testWorkflowEnvironmentImpl) GetVersion(changeID string, minSupported, maxSupported Version) (retVersion Version) {
  1908  	if mockVersion, ok := env.getMockedVersion(changeID, changeID, minSupported, maxSupported); ok {
  1909  		// GetVersion for changeID is mocked
  1910  		env.UpsertSearchAttributes(createSearchAttributesForChangeVersion(changeID, mockVersion, env.changeVersions))
  1911  		env.changeVersions[changeID] = mockVersion
  1912  		return mockVersion
  1913  	}
  1914  	if mockVersion, ok := env.getMockedVersion(mock.Anything, changeID, minSupported, maxSupported); ok {
  1915  		// GetVersion is mocked with any changeID.
  1916  		env.UpsertSearchAttributes(createSearchAttributesForChangeVersion(changeID, mockVersion, env.changeVersions))
  1917  		env.changeVersions[changeID] = mockVersion
  1918  		return mockVersion
  1919  	}
  1920  
  1921  	// no mock setup, so call regular path
  1922  	if version, ok := env.changeVersions[changeID]; ok {
  1923  		validateVersion(changeID, version, minSupported, maxSupported)
  1924  		return version
  1925  	}
  1926  	env.UpsertSearchAttributes(createSearchAttributesForChangeVersion(changeID, maxSupported, env.changeVersions))
  1927  	env.changeVersions[changeID] = maxSupported
  1928  	return maxSupported
  1929  }
  1930  
  1931  func (env *testWorkflowEnvironmentImpl) getMockedVersion(mockedChangeID, changeID string, minSupported, maxSupported Version) (Version, bool) {
  1932  	mockMethod := getMockMethodForGetVersion(mockedChangeID)
  1933  	if _, ok := env.expectedMockCalls[mockMethod]; !ok {
  1934  		// mock not found
  1935  		return DefaultVersion, false
  1936  	}
  1937  
  1938  	args := []interface{}{changeID, minSupported, maxSupported}
  1939  	// below call will panic if mock is not properly setup.
  1940  	mockRet := env.mock.MethodCalled(mockMethod, args...)
  1941  	m := &mockWrapper{name: mockMethodForGetVersion, fn: mockFnGetVersion}
  1942  	if mockFn := m.getMockFn(mockRet); mockFn != nil {
  1943  		executor := &activityExecutor{name: mockMethodForGetVersion, fn: mockFn}
  1944  		reflectValues := executor.executeWithActualArgsWithoutParseResult(nil, args)
  1945  		if len(reflectValues) != 1 || !reflect.TypeOf(reflectValues[0].Interface()).AssignableTo(reflect.TypeOf(DefaultVersion)) {
  1946  			panic(fmt.Sprintf("mock of GetVersion has incorrect return type, expected workflow.Version, but actual is %T (%v)",
  1947  				reflectValues[0].Interface(), reflectValues[0].Interface()))
  1948  		}
  1949  		return reflectValues[0].Interface().(Version), true
  1950  	}
  1951  
  1952  	if len(mockRet) != 1 || !reflect.TypeOf(mockRet[0]).AssignableTo(reflect.TypeOf(DefaultVersion)) {
  1953  		panic(fmt.Sprintf("mock of GetVersion has incorrect return type, expected workflow.Version, but actual is %T (%v)",
  1954  			mockRet[0], mockRet[0]))
  1955  	}
  1956  	return mockRet[0].(Version), true
  1957  }
  1958  
  1959  func getMockMethodForGetVersion(changeID string) string {
  1960  	return fmt.Sprintf("%v_%v", mockMethodForGetVersion, changeID)
  1961  }
  1962  
  1963  func (env *testWorkflowEnvironmentImpl) UpsertSearchAttributes(attributes map[string]interface{}) error {
  1964  	mockMethod := mockMethodForUpsertSearchAttributes
  1965  	if _, ok := env.expectedMockCalls[mockMethod]; ok {
  1966  		// mock found, check if return is error
  1967  		args := []interface{}{attributes}
  1968  		mockRet := env.mock.MethodCalled(mockMethod, args...)
  1969  		if len(mockRet) > 1 {
  1970  			panic(fmt.Sprintf("mock of UpsertSearchAttributes should return only one error"))
  1971  		}
  1972  		if len(mockRet) == 1 && mockRet[0] != nil {
  1973  			return mockRet[0].(error)
  1974  		}
  1975  	}
  1976  
  1977  	attr, err := validateAndSerializeSearchAttributes(attributes)
  1978  	env.workflowInfo.SearchAttributes = mergeSearchAttributes(env.workflowInfo.SearchAttributes, attr)
  1979  	return err
  1980  }
  1981  
  1982  func (env *testWorkflowEnvironmentImpl) MutableSideEffect(id string, f func() interface{}, equals func(a, b interface{}) bool) Value {
  1983  	return newEncodedValue(env.encodeValue(f()), env.GetDataConverter())
  1984  }
  1985  
  1986  func (env *testWorkflowEnvironmentImpl) AddSession(sessionInfo *SessionInfo) {
  1987  	env.openSessions[sessionInfo.SessionID] = sessionInfo
  1988  }
  1989  
  1990  func (env *testWorkflowEnvironmentImpl) RemoveSession(sessionID string) {
  1991  	delete(env.openSessions, sessionID)
  1992  }
  1993  
  1994  func (env *testWorkflowEnvironmentImpl) encodeValue(value interface{}) []byte {
  1995  	blob, err := env.GetDataConverter().ToData(value)
  1996  	if err != nil {
  1997  		panic(err)
  1998  	}
  1999  	return blob
  2000  }
  2001  
  2002  func (env *testWorkflowEnvironmentImpl) nextID() int {
  2003  	activityID := env.counterID
  2004  	env.counterID++
  2005  	return activityID
  2006  }
  2007  
  2008  func getStringID(intID int) string {
  2009  	return fmt.Sprintf("%d", intID)
  2010  }
  2011  
  2012  func (env *testWorkflowEnvironmentImpl) getActivityInfo(activityID, activityType string) *ActivityInfo {
  2013  	return &ActivityInfo{
  2014  		ActivityID:        activityID,
  2015  		ActivityType:      ActivityType{Name: activityType},
  2016  		TaskToken:         []byte(activityID),
  2017  		WorkflowExecution: env.workflowInfo.WorkflowExecution,
  2018  	}
  2019  }
  2020  
  2021  func (env *testWorkflowEnvironmentImpl) cancelWorkflow(callback resultHandler) {
  2022  	env.postCallback(func() {
  2023  		// RequestCancelWorkflow needs to be run in main thread
  2024  		env.RequestCancelExternalWorkflow(
  2025  			env.workflowInfo.Domain,
  2026  			env.workflowInfo.WorkflowExecution.ID,
  2027  			env.workflowInfo.WorkflowExecution.RunID,
  2028  			callback,
  2029  		)
  2030  	}, true)
  2031  }
  2032  
  2033  func (env *testWorkflowEnvironmentImpl) signalWorkflow(name string, input interface{}, startDecisionTask bool) {
  2034  	data, err := encodeArg(env.GetDataConverter(), input)
  2035  	if err != nil {
  2036  		panic(err)
  2037  	}
  2038  	env.postCallback(func() {
  2039  		env.signalHandler(name, data)
  2040  	}, startDecisionTask)
  2041  }
  2042  
  2043  func (env *testWorkflowEnvironmentImpl) signalWorkflowByID(workflowID, signalName string, input interface{}) error {
  2044  	data, err := encodeArg(env.GetDataConverter(), input)
  2045  	if err != nil {
  2046  		panic(err)
  2047  	}
  2048  
  2049  	if workflowHandle, ok := env.runningWorkflows[workflowID]; ok {
  2050  		if workflowHandle.handled {
  2051  			return &shared.WorkflowExecutionAlreadyCompletedError{Message: fmt.Sprintf("Workflow %v already completed", workflowID)}
  2052  		}
  2053  		workflowHandle.env.postCallback(func() {
  2054  			workflowHandle.env.signalHandler(signalName, data)
  2055  		}, true)
  2056  		return nil
  2057  	}
  2058  
  2059  	return &shared.EntityNotExistsError{Message: fmt.Sprintf("Workflow %v not exists", workflowID)}
  2060  }
  2061  
  2062  func (env *testWorkflowEnvironmentImpl) queryWorkflow(queryType string, args ...interface{}) (Value, error) {
  2063  	data, err := encodeArgs(env.GetDataConverter(), args)
  2064  	if err != nil {
  2065  		return nil, err
  2066  	}
  2067  	blob, err := env.queryHandler(queryType, data)
  2068  	if err != nil {
  2069  		return nil, err
  2070  	}
  2071  	return newEncodedValue(blob, env.GetDataConverter()), nil
  2072  }
  2073  
  2074  func (env *testWorkflowEnvironmentImpl) getMockRunFn(callWrapper *MockCallWrapper) func(args mock.Arguments) {
  2075  	env.locker.Lock()
  2076  	defer env.locker.Unlock()
  2077  
  2078  	env.expectedMockCalls[callWrapper.call.Method] = struct{}{}
  2079  	return func(args mock.Arguments) {
  2080  		env.runBeforeMockCallReturns(callWrapper, args)
  2081  	}
  2082  }
  2083  
  2084  func (env *testWorkflowEnvironmentImpl) setLastCompletionResult(result interface{}) {
  2085  	data, err := encodeArg(env.GetDataConverter(), result)
  2086  	if err != nil {
  2087  		panic(err)
  2088  	}
  2089  	env.workflowInfo.lastCompletionResult = data
  2090  }
  2091  
  2092  func (env *testWorkflowEnvironmentImpl) setHeartbeatDetails(details interface{}) {
  2093  	data, err := encodeArg(env.GetDataConverter(), details)
  2094  	if err != nil {
  2095  		panic(err)
  2096  	}
  2097  	env.heartbeatDetails = data
  2098  }
  2099  
  2100  func (env *testWorkflowEnvironmentImpl) GetRegistry() *registry {
  2101  	return env.registry
  2102  }
  2103  
  2104  func (env *testWorkflowEnvironmentImpl) GetWorkflowInterceptors() []WorkflowInterceptorFactory {
  2105  	return env.workflowInterceptors
  2106  }
  2107  
  2108  func newTestSessionEnvironment(testWorkflowEnvironment *testWorkflowEnvironmentImpl,
  2109  	params *workerExecutionParameters, concurrentSessionExecutionSize int) *testSessionEnvironmentImpl {
  2110  	resourceID := params.SessionResourceID
  2111  	if resourceID == "" {
  2112  		resourceID = "testResourceID"
  2113  	}
  2114  	if concurrentSessionExecutionSize == 0 {
  2115  		concurrentSessionExecutionSize = defaultMaxConcurrentSessionExecutionSize
  2116  	}
  2117  
  2118  	return &testSessionEnvironmentImpl{
  2119  		sessionEnvironmentImpl:  newSessionEnvironment(resourceID, concurrentSessionExecutionSize).(*sessionEnvironmentImpl),
  2120  		testWorkflowEnvironment: testWorkflowEnvironment,
  2121  	}
  2122  }
  2123  
  2124  func (t *testSessionEnvironmentImpl) SignalCreationResponse(ctx context.Context, sessionID string) error {
  2125  	t.testWorkflowEnvironment.signalWorkflow(sessionID, t.sessionEnvironmentImpl.getCreationResponse(), true)
  2126  	return nil
  2127  }
  2128  
  2129  // function signature for mock SignalExternalWorkflow
  2130  func mockFnSignalExternalWorkflow(domainName, workflowID, runID, signalName string, arg interface{}) error {
  2131  	return nil
  2132  }
  2133  
  2134  // function signature for mock RequestCancelExternalWorkflow
  2135  func mockFnRequestCancelExternalWorkflow(domainName, workflowID, runID string) error {
  2136  	return nil
  2137  }
  2138  
  2139  // function signature for mock GetVersion
  2140  func mockFnGetVersion(changeID string, minSupported, maxSupported Version) Version {
  2141  	return DefaultVersion
  2142  }
  2143  
  2144  type testReporter struct {
  2145  	logger *zap.Logger
  2146  }
  2147  
  2148  func (t *testReporter) Errorf(format string, args ...interface{}) {
  2149  	t.logger.Error(fmt.Sprintf(format, args...))
  2150  }
  2151  
  2152  func (t *testReporter) Fatalf(format string, args ...interface{}) {
  2153  	t.logger.Fatal(fmt.Sprintf(format, args...))
  2154  }