fix flakey unit test (#238)

* fix flakey unit test

* define activityProvider

Author: yiminc

internal_task_handlers.go

@@ -85,15 +85,17 @@ type (
 		hostEnv               *hostEnvImpl
 	}
 
+	activityProvider func(name string) activity
 	// activityTaskHandlerImpl is the implementation of ActivityTaskHandler
 	activityTaskHandlerImpl struct {
-		taskListName string
-		identity     string
-		service      m.TChanWorkflowService
-		metricsScope tally.Scope
-		logger       *zap.Logger
-		userContext  context.Context
-		hostEnv      *hostEnvImpl
+		taskListName     string
+		identity         string
+		service          m.TChanWorkflowService
+		metricsScope     tally.Scope
+		logger           *zap.Logger
+		userContext      context.Context
+		hostEnv          *hostEnvImpl
+		activityProvider activityProvider
 	}
 
 	// history wrapper method to help information about events.
@@ -809,15 +811,25 @@ func newActivityTaskHandler(
 	service m.TChanWorkflowService,
 	params workerExecutionParameters,
 	env *hostEnvImpl,
+) ActivityTaskHandler {
+	return newActivityTaskHandlerWithCustomProvider(service, params, env, nil)
+}
+
+func newActivityTaskHandlerWithCustomProvider(
+	service m.TChanWorkflowService,
+	params workerExecutionParameters,
+	env *hostEnvImpl,
+	activityProvider activityProvider,
 ) ActivityTaskHandler {
 	return &activityTaskHandlerImpl{
-		taskListName: params.TaskList,
-		identity:     params.Identity,
-		service:      service,
-		logger:       params.Logger,
-		metricsScope: params.MetricsScope,
-		userContext:  params.UserContext,
-		hostEnv:      env,
+		taskListName:     params.TaskList,
+		identity:         params.Identity,
+		service:          service,
+		logger:           params.Logger,
+		metricsScope:     params.MetricsScope,
+		userContext:      params.UserContext,
+		hostEnv:          env,
+		activityProvider: activityProvider,
 	}
 }
 
@@ -958,10 +970,10 @@ func (ath *activityTaskHandlerImpl) Execute(t *s.PollForActivityTaskResponse) (r
 	defer invoker.Close()
 	ctx := WithActivityTask(canCtx, t, invoker, ath.logger, ath.metricsScope)
 	activityType := *t.GetActivityType()
-	activityImplementation, ok := ath.hostEnv.getActivity(activityType.GetName())
-	if !ok {
+	activityImplementation := ath.getActivity(activityType.GetName())
+	if activityImplementation == nil {
 		// Couldn't find the activity implementation.
-		return nil, fmt.Errorf("Unable to find activityType=%v", activityType.GetName())
+		return nil, fmt.Errorf("unable to find activityType=%v", activityType.GetName())
 	}
 
 	// panic handler
@@ -999,6 +1011,18 @@ func (ath *activityTaskHandlerImpl) Execute(t *s.PollForActivityTaskResponse) (r
 	return convertActivityResultToRespondRequest(ath.identity, t.TaskToken, output, err), nil
 }
 
+func (ath *activityTaskHandlerImpl) getActivity(name string) activity {
+	if ath.activityProvider != nil {
+		return ath.activityProvider(name)
+	}
+
+	if a, ok := ath.hostEnv.getActivity(name); ok {
+		return a
+	}
+
+	return nil
+}
+
 func createNewDecision(decisionType s.DecisionType) *s.Decision {
 	return &s.Decision{
 		DecisionType: common.DecisionTypePtr(decisionType),

internal_workflow_testsuite.go

@@ -945,42 +945,29 @@ func (env *testWorkflowEnvironmentImpl) newTestActivityTaskHandler(taskList stri
 	}
 	ensureRequiredParams(&params)
 
-	for fnName, tasklistActivity := range env.taskListSpecificActivities {
-		if _, ok := tasklistActivity.taskLists[taskList]; ok {
-			ae := &activityExecutor{name: fnName, fn: tasklistActivity.fn}
-			getHostEnvironment().addActivity(
-				fnName,
-				&activityExecutorWrapper{activityExecutor: ae, env: env},
-			)
-		}
-	}
-
 	if len(getHostEnvironment().getRegisteredActivities()) == 0 {
 		panic(fmt.Sprintf("no activity is registered for tasklist '%v'", taskList))
 	}
 
-	for _, a := range getHostEnvironment().getRegisteredActivities() {
-		fnName := a.ActivityType().Name
-		if _, ok := env.taskListSpecificActivities[fnName]; ok {
-			// activity is registered to a specific taskList, so ignore it from the global registered activities.
-			continue
+	getActivity := func(name string) activity {
+		tlsa, ok := env.taskListSpecificActivities[name]
+		if ok {
+			_, ok := tlsa.taskLists[taskList]
+			if !ok {
+				// activity are bind to specific task list but not to current task list
+				return nil
+			}
 		}
-		var ae *activityExecutor
-		switch v := a.(type) {
-		case *activityExecutor:
-			ae = v
-		case *activityExecutorWrapper:
-			ae = v.activityExecutor
-		default:
-			ae = &activityExecutor{name: v.ActivityType().Name, fn: v.GetFunction()}
+
+		activity, ok := getHostEnvironment().getActivity(name)
+		if !ok {
+			return nil
 		}
-		getHostEnvironment().addActivity(
-			fnName,
-			&activityExecutorWrapper{activityExecutor: ae, env: env},
-		)
+		ae := &activityExecutor{name: activity.ActivityType().Name, fn: activity.GetFunction()}
+		return &activityExecutorWrapper{activityExecutor: ae, env: env}
 	}
 
-	taskHandler := newActivityTaskHandler(env.service, params, getHostEnvironment())
+	taskHandler := newActivityTaskHandlerWithCustomProvider(env.service, params, getHostEnvironment(), getActivity)
 	return taskHandler
 }
 

internal_workflow_testsuite_test.go

@@ -92,10 +92,11 @@ func (s *WorkflowTestSuiteUnitTest) Test_ActivityMockValues() {
 }
 
 func (s *WorkflowTestSuiteUnitTest) Test_OnActivityStartedListener() {
+	runCount := 100
 	workflowFn := func(ctx Context) error {
 		ctx = WithActivityOptions(ctx, s.activityOptions)
 
-		for i := 1; i <= 3; i++ {
+		for i := 1; i <= runCount; i++ {
 			err := ExecuteActivity(ctx, testActivityHello, fmt.Sprintf("msg%d", i)).Get(ctx, nil)
 			if err != nil {
 				return err
@@ -113,10 +114,9 @@ func (s *WorkflowTestSuiteUnitTest) Test_OnActivityStartedListener() {
 		s.NoError(args.Get(&input))
 		activityCalls = append(activityCalls, fmt.Sprintf("%s:%s", activityInfo.ActivityType.Name, input))
 	})
-	expectedCalls := []string{
-		"testActivityHello:msg1",
-		"testActivityHello:msg2",
-		"testActivityHello:msg3",
+	expectedCalls := []string{}
+	for i := 1; i <= runCount; i++ {
+		expectedCalls = append(expectedCalls, fmt.Sprintf("testActivityHello:msg%v", i))
 	}
 
 	env.ExecuteWorkflow(workflowFn)