go.uber.org/cadence@v1.2.9/internal/internal_workflow_testsuite_test.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  	"encoding/json"
    27  	"errors"
    28  	"fmt"
    29  	"runtime"
    30  	"sync"
    31  	"testing"
    32  	"time"
    33  
    34  	"github.com/stretchr/testify/assert"
    35  	"github.com/stretchr/testify/mock"
    36  	"github.com/stretchr/testify/require"
    37  	"github.com/stretchr/testify/suite"
    38  	"go.uber.org/zap"
    39  	"go.uber.org/zap/zaptest"
    40  
    41  	"go.uber.org/cadence/.gen/go/shared"
    42  	"go.uber.org/cadence/internal/common"
    43  )
    44  
    45  type WorkflowTestSuiteUnitTest struct {
    46  	suite.Suite
    47  	WorkflowTestSuite
    48  	activityOptions      ActivityOptions
    49  	localActivityOptions LocalActivityOptions
    50  }
    51  
    52  type testContextKey string
    53  
    54  func (s *WorkflowTestSuiteUnitTest) SetupSuite() {
    55  	s.activityOptions = ActivityOptions{
    56  		ScheduleToStartTimeout: time.Minute,
    57  		StartToCloseTimeout:    time.Minute,
    58  		HeartbeatTimeout:       20 * time.Second,
    59  	}
    60  	s.localActivityOptions = LocalActivityOptions{
    61  		ScheduleToCloseTimeout: time.Second * 3,
    62  	}
    63  	s.header = &shared.Header{
    64  		Fields: map[string][]byte{"test": []byte("test-data")},
    65  	}
    66  	s.ctxProps = []ContextPropagator{NewStringMapPropagator([]string{"test"})}
    67  }
    68  
    69  func (s *WorkflowTestSuiteUnitTest) SetupTest() {
    70  	s.SetLogger(zaptest.NewLogger(s.T()))
    71  }
    72  
    73  func TestUnitTestSuite(t *testing.T) {
    74  	suite.Run(t, new(WorkflowTestSuiteUnitTest))
    75  }
    76  
    77  func (s *WorkflowTestSuiteUnitTest) Test_ActivityMockFunction() {
    78  	mockActivity := func(ctx context.Context, msg string) (string, error) {
    79  		return "mock_" + msg, nil
    80  	}
    81  
    82  	env := s.NewTestWorkflowEnvironment()
    83  	env.RegisterActivity(testActivityHello)
    84  	env.OnActivity(testActivityHello, mock.Anything, mock.Anything).Return(mockActivity).Once()
    85  	env.RegisterWorkflow(testWorkflowHello)
    86  	env.ExecuteWorkflow(testWorkflowHello)
    87  
    88  	s.True(env.IsWorkflowCompleted())
    89  	s.NoError(env.GetWorkflowError())
    90  	var result string
    91  	env.GetWorkflowResult(&result)
    92  	s.Equal("mock_world", result)
    93  	env.AssertExpectations(s.T())
    94  }
    95  
    96  func (s *WorkflowTestSuiteUnitTest) Test_ActivityMockFunction_WithDataConverter() {
    97  	mockActivity := func(ctx context.Context, msg string) (string, error) {
    98  		return "mock_" + msg, nil
    99  	}
   100  
   101  	workflowFn := func(ctx Context) (string, error) {
   102  		ao := ActivityOptions{
   103  			ScheduleToStartTimeout: time.Minute,
   104  			StartToCloseTimeout:    time.Minute,
   105  			HeartbeatTimeout:       20 * time.Second,
   106  		}
   107  		ctx = WithActivityOptions(ctx, ao)
   108  
   109  		var result string
   110  		ctx = WithDataConverter(ctx, newTestDataConverter())
   111  		err := ExecuteActivity(ctx, testActivityHello, "world").Get(ctx, &result)
   112  		if err != nil {
   113  			return "", err
   114  		}
   115  
   116  		var result1 string
   117  		ctx1 := WithDataConverter(ctx, getDefaultDataConverter()) // use another converter to run activity
   118  		err1 := ExecuteActivity(ctx1, testActivityHello, "world1").Get(ctx1, &result1)
   119  		if err1 != nil {
   120  			return "", err1
   121  		}
   122  		return result + "," + result1, nil
   123  	}
   124  
   125  	env := s.NewTestWorkflowEnvironment()
   126  	env.RegisterWorkflow(workflowFn)
   127  	env.RegisterActivity(testActivityHello)
   128  
   129  	env.SetWorkerOptions(WorkerOptions{DataConverter: newTestDataConverter()})
   130  	env.OnActivity(testActivityHello, mock.Anything, mock.Anything).Return(mockActivity).Twice()
   131  
   132  	env.ExecuteWorkflow(workflowFn)
   133  
   134  	s.True(env.IsWorkflowCompleted())
   135  	s.NoError(env.GetWorkflowError())
   136  	var result string
   137  	env.GetWorkflowResult(&result)
   138  	s.Equal("mock_world,mock_world1", result)
   139  	env.AssertExpectations(s.T())
   140  }
   141  
   142  func (s *WorkflowTestSuiteUnitTest) Test_ActivityMockValues() {
   143  	env := s.NewTestWorkflowEnvironment()
   144  	env.RegisterWorkflow(testWorkflowHello)
   145  	env.RegisterActivity(testActivityHello)
   146  
   147  	env.OnActivity(testActivityHello, mock.Anything, mock.Anything).Return("mock_value", nil).Once()
   148  	env.ExecuteWorkflow(testWorkflowHello)
   149  
   150  	s.True(env.IsWorkflowCompleted())
   151  	s.NoError(env.GetWorkflowError())
   152  	var result string
   153  	env.GetWorkflowResult(&result)
   154  	s.Equal("mock_value", result)
   155  	env.AssertExpectations(s.T())
   156  }
   157  
   158  func (s *WorkflowTestSuiteUnitTest) Test_OnActivityStartedListener() {
   159  	runCount := 100
   160  	workflowFn := func(ctx Context) error {
   161  		ctx = WithActivityOptions(ctx, s.activityOptions)
   162  
   163  		for i := 1; i <= runCount; i++ {
   164  			err := ExecuteActivity(ctx, testActivityHello, fmt.Sprintf("msg%d", i)).Get(ctx, nil)
   165  			if err != nil {
   166  				return err
   167  			}
   168  		}
   169  		return nil
   170  	} // END of workflow code
   171  
   172  	env := s.NewTestWorkflowEnvironment()
   173  	env.RegisterWorkflow(workflowFn)
   174  	env.RegisterActivityWithOptions(testActivityHello, RegisterActivityOptions{Name: "testActivityHello"})
   175  
   176  	var activityCalls []string
   177  	env.SetOnActivityStartedListener(func(activityInfo *ActivityInfo, ctx context.Context, args Values) {
   178  		var input string
   179  		s.NoError(args.Get(&input))
   180  		activityCalls = append(activityCalls, fmt.Sprintf("%s:%s", activityInfo.ActivityType.Name, input))
   181  	})
   182  	var expectedCalls []string
   183  	for i := 1; i <= runCount; i++ {
   184  		expectedCalls = append(expectedCalls, fmt.Sprintf("testActivityHello:msg%v", i))
   185  	}
   186  
   187  	env.ExecuteWorkflow(workflowFn)
   188  	s.True(env.IsWorkflowCompleted())
   189  	s.NoError(env.GetWorkflowError())
   190  	s.Equal(expectedCalls, activityCalls)
   191  }
   192  
   193  func (s *WorkflowTestSuiteUnitTest) Test_TimerWorkflow_ClockAutoFastForward() {
   194  	var firedTimerRecord []string
   195  	workflowFn := func(ctx Context) error {
   196  		t1 := NewTimer(ctx, time.Second*5)
   197  		t2 := NewTimer(ctx, time.Second*1)
   198  		t3 := NewTimer(ctx, time.Second*2)
   199  		t4 := NewTimer(ctx, time.Second*5)
   200  
   201  		selector := NewSelector(ctx)
   202  		selector.AddFuture(t1, func(f Future) {
   203  			firedTimerRecord = append(firedTimerRecord, "t1")
   204  		}).AddFuture(t2, func(f Future) {
   205  			firedTimerRecord = append(firedTimerRecord, "t2")
   206  		}).AddFuture(t3, func(f Future) {
   207  			firedTimerRecord = append(firedTimerRecord, "t3")
   208  		}).AddFuture(t4, func(f Future) {
   209  			firedTimerRecord = append(firedTimerRecord, "t4")
   210  		})
   211  
   212  		selector.Select(ctx)
   213  		selector.Select(ctx)
   214  		selector.Select(ctx)
   215  		selector.Select(ctx)
   216  
   217  		return nil
   218  	}
   219  
   220  	env := s.NewTestWorkflowEnvironment()
   221  	env.RegisterWorkflow(workflowFn)
   222  	env.ExecuteWorkflow(workflowFn)
   223  
   224  	s.True(env.IsWorkflowCompleted())
   225  	s.NoError(env.GetWorkflowError())
   226  	s.Equal([]string{"t2", "t3", "t1", "t4"}, firedTimerRecord)
   227  }
   228  
   229  func (s *WorkflowTestSuiteUnitTest) Test_WorkflowAutoForwardClock() {
   230  	workflowFn := func(ctx Context) (string, error) {
   231  		// Schedule a timer with long duration. In this test, we won't actually wait for that long, because the test suite
   232  		// will auto forward clock when workflow is blocked and there is no running activity.
   233  		f1 := NewTimer(ctx, time.Minute)
   234  
   235  		ctx = WithActivityOptions(ctx, s.activityOptions)
   236  		// Execute activity that returns immediately, once the activity returns, the workflow will be blocked on timer
   237  		// and the test suite will auto forward clock to fire the timer.
   238  		f2 := ExecuteActivity(ctx, testActivityHello, "controlled_execution")
   239  
   240  		timerErr := f1.Get(ctx, nil) // wait until timer fires
   241  		if timerErr != nil {
   242  			return "", timerErr
   243  		}
   244  
   245  		if !f2.IsReady() {
   246  			return "", errors.New("activity is not completed when timer fired")
   247  		}
   248  
   249  		var activityResult string
   250  		activityErr := f2.Get(ctx, &activityResult)
   251  		if activityErr != nil {
   252  			return "", activityErr
   253  		}
   254  
   255  		return activityResult, nil
   256  	} // END of workflow code
   257  
   258  	env := s.NewTestWorkflowEnvironment()
   259  	env.RegisterWorkflow(workflowFn)
   260  	env.RegisterActivity(testActivityHello)
   261  	env.ExecuteWorkflow(workflowFn)
   262  
   263  	s.True(env.IsWorkflowCompleted())
   264  	s.NoError(env.GetWorkflowError())
   265  	var result string
   266  	env.GetWorkflowResult(&result)
   267  	s.Equal("hello_controlled_execution", result)
   268  }
   269  
   270  func (s *WorkflowTestSuiteUnitTest) Test_WorkflowMixedClock() {
   271  	workflowFn := func(ctx Context) (string, error) {
   272  		// Schedule a long timer.
   273  		t1 := NewTimer(ctx, time.Minute)
   274  
   275  		ctx = WithActivityOptions(ctx, s.activityOptions)
   276  		// Schedule 2 activities, one returns immediately, the other will take 10s to run.
   277  		f1 := ExecuteActivity(ctx, testActivityHello, "controlled_execution")
   278  		f2Ctx, f2Cancel := WithCancel(ctx)
   279  		f2 := ExecuteActivity(f2Ctx, testActivityHeartbeat, time.Second*10)
   280  
   281  		err := f1.Get(ctx, nil)
   282  		if err != nil {
   283  			return "", err
   284  		}
   285  
   286  		// Schedule a short timer after f1 completed.
   287  		t2 := NewTimer(ctx, time.Millisecond)
   288  		NewSelector(ctx).AddFuture(t2, func(f Future) {
   289  			// when t2 fires, we would cancel f2
   290  			if !f2.IsReady() {
   291  				f2Cancel()
   292  			}
   293  		}).Select(ctx) // wait until t2 fires
   294  
   295  		t1.Get(ctx, nil) // wait for the long timer to fire.
   296  
   297  		return "expected", nil
   298  	} // END of workflow code
   299  
   300  	env := s.NewTestWorkflowEnvironment()
   301  	env.RegisterWorkflow(workflowFn)
   302  	env.RegisterActivity(testActivityHeartbeat)
   303  	env.RegisterActivity(testActivityHello)
   304  
   305  	env.ExecuteWorkflow(workflowFn)
   306  
   307  	s.True(env.IsWorkflowCompleted())
   308  	s.NoError(env.GetWorkflowError())
   309  	var result string
   310  	env.GetWorkflowResult(&result)
   311  	s.Equal("expected", result)
   312  }
   313  
   314  func (s *WorkflowTestSuiteUnitTest) Test_WorkflowActivityCancellation() {
   315  	workflowFn := func(ctx Context) error {
   316  		ctx = WithActivityOptions(ctx, s.activityOptions)
   317  
   318  		ctx, cancelHandler := WithCancel(ctx)
   319  		f1 := ExecuteActivity(ctx, testActivityHeartbeat, "fast", time.Millisecond) // fast activity
   320  		f2 := ExecuteActivity(ctx, testActivityHeartbeat, "slow", time.Second*3)    // slow activity
   321  
   322  		NewSelector(ctx).AddFuture(f1, func(f Future) {
   323  			cancelHandler()
   324  		}).AddFuture(f2, func(f Future) {
   325  			cancelHandler()
   326  		}).Select(ctx)
   327  
   328  		err := f2.Get(ctx, nil) // verify slow activity is cancelled
   329  		if _, ok := err.(*CanceledError); !ok {
   330  			return err
   331  		}
   332  		return nil
   333  	}
   334  
   335  	env := s.NewTestWorkflowEnvironment()
   336  	env.RegisterWorkflow(workflowFn)
   337  	env.RegisterActivity(testActivityHeartbeat)
   338  	activityMap := make(map[string]string) // msg -> activityID
   339  	var completedActivityID, cancelledActivityID string
   340  	env.SetOnActivityStartedListener(func(activityInfo *ActivityInfo, ctx context.Context, args Values) {
   341  		var msg string
   342  		s.NoError(args.Get(&msg))
   343  		activityMap[msg] = activityInfo.ActivityID
   344  	})
   345  	env.SetOnActivityCompletedListener(func(activityInfo *ActivityInfo, result Value, err error) {
   346  		completedActivityID = activityInfo.ActivityID
   347  	})
   348  	env.SetOnActivityCanceledListener(func(activityInfo *ActivityInfo) {
   349  		cancelledActivityID = activityInfo.ActivityID
   350  	})
   351  	env.ExecuteWorkflow(workflowFn)
   352  
   353  	s.True(env.IsWorkflowCompleted())
   354  	s.NoError(env.GetWorkflowError())
   355  	s.Equal(activityMap["fast"], completedActivityID)
   356  	s.Equal(activityMap["slow"], cancelledActivityID)
   357  }
   358  
   359  func (s *WorkflowTestSuiteUnitTest) Test_ActivityWithUserContext() {
   360  	testKey, testValue := testContextKey("test_key"), "test_value"
   361  	userCtx := context.WithValue(context.Background(), testKey, testValue)
   362  	workerOptions := WorkerOptions{}
   363  	workerOptions.BackgroundActivityContext = userCtx
   364  
   365  	// inline activity using value passing through user context.
   366  	activityWithUserContext := func(ctx context.Context, keyName testContextKey) (string, error) {
   367  		value := ctx.Value(keyName)
   368  		if value != nil {
   369  			return value.(string), nil
   370  		}
   371  		return "", errors.New("value not found from ctx")
   372  	}
   373  
   374  	env := s.NewTestActivityEnvironment()
   375  	env.RegisterActivity(activityWithUserContext)
   376  	env.SetWorkerOptions(workerOptions)
   377  	blob, err := env.ExecuteActivity(activityWithUserContext, testKey)
   378  	s.NoError(err)
   379  	var value string
   380  	blob.Get(&value)
   381  	s.Equal(testValue, value)
   382  }
   383  
   384  func (s *WorkflowTestSuiteUnitTest) Test_ActivityWithHeaderContext() {
   385  	workerOptions := WorkerOptions{
   386  		ContextPropagators: []ContextPropagator{NewStringMapPropagator([]string{testHeader})},
   387  	}
   388  
   389  	// inline activity using value passing through user context.
   390  	activityWithUserContext := func(ctx context.Context) (string, error) {
   391  		value := ctx.Value(contextKey(testHeader))
   392  		if val, ok := value.(string); ok {
   393  			return val, nil
   394  		}
   395  		return "", errors.New("value not found from ctx")
   396  	}
   397  
   398  	s.SetHeader(&shared.Header{
   399  		Fields: map[string][]byte{
   400  			testHeader: []byte("test-data"),
   401  		},
   402  	})
   403  
   404  	env := s.NewTestActivityEnvironment()
   405  	env.RegisterActivity(activityWithUserContext)
   406  	env.SetWorkerOptions(workerOptions)
   407  	blob, err := env.ExecuteActivity(activityWithUserContext)
   408  	s.NoError(err)
   409  	var value string
   410  	blob.Get(&value)
   411  	s.Equal("test-data", value)
   412  }
   413  
   414  func (s *WorkflowTestSuiteUnitTest) Test_CompleteActivity() {
   415  	env := s.NewTestWorkflowEnvironment()
   416  	var activityInfo ActivityInfo
   417  	mockActivity := func(ctx context.Context, msg string) (string, error) {
   418  		activityInfo = GetActivityInfo(ctx)
   419  		env.RegisterDelayedCallback(func() {
   420  			err := env.CompleteActivity(activityInfo.TaskToken, "async_complete", nil)
   421  			s.NoError(err)
   422  		}, time.Minute)
   423  		return "", ErrActivityResultPending
   424  	}
   425  
   426  	env.RegisterWorkflow(testWorkflowHello)
   427  	env.RegisterActivity(testActivityHello)
   428  	env.OnActivity(testActivityHello, mock.Anything, mock.Anything).Return(mockActivity).Once()
   429  
   430  	env.ExecuteWorkflow(testWorkflowHello)
   431  
   432  	s.True(env.IsWorkflowCompleted())
   433  	s.NoError(env.GetWorkflowError())
   434  	env.AssertExpectations(s.T())
   435  
   436  	var result string
   437  	env.GetWorkflowResult(&result)
   438  	s.Equal("async_complete", result)
   439  }
   440  
   441  func (s *WorkflowTestSuiteUnitTest) Test_ActivityReturnsErrActivityResultPending() {
   442  	env := s.NewTestActivityEnvironment()
   443  	activityFn := func(ctx context.Context) (string, error) {
   444  		return "", ErrActivityResultPending
   445  	}
   446  	env.RegisterActivity(activityFn)
   447  	_, err := env.ExecuteActivity(activityFn)
   448  	s.Equal(ErrActivityResultPending, err)
   449  }
   450  
   451  func (s *WorkflowTestSuiteUnitTest) Test_WorkflowCancellation() {
   452  	workflowFn := func(ctx Context) error {
   453  		ctx = WithActivityOptions(ctx, s.activityOptions)
   454  		f := ExecuteActivity(ctx, testActivityHeartbeat, "msg1", time.Second*10)
   455  		err := f.Get(ctx, nil) // wait for result
   456  		return err
   457  	}
   458  
   459  	env := s.NewTestWorkflowEnvironment()
   460  	// Register a delayed callback using workflow timer internally. The callback will be called when workflow clock passed
   461  	// by the specified delay duration. The test suite enables the auto clock forwarding when workflow is blocked and no
   462  	// activities are running. However, if there are running activities, test suite's workflow clock will move forward at
   463  	// the real wall clock pace. In this test case, the activity is configured to run for 10s, so the workflow will be
   464  	// blocked. But after 1ms, the registered callback will be invoked, which will request to cancel the workflow.
   465  	env.RegisterDelayedCallback(func() {
   466  		env.CancelWorkflow()
   467  	}, time.Millisecond)
   468  
   469  	env.RegisterWorkflow(workflowFn)
   470  	env.RegisterActivity(testActivityHeartbeat)
   471  
   472  	env.ExecuteWorkflow(workflowFn)
   473  
   474  	s.True(env.IsWorkflowCompleted())
   475  	s.NotNil(env.GetWorkflowError())
   476  	_, ok := env.GetWorkflowError().(*CanceledError)
   477  	s.True(ok)
   478  }
   479  
   480  func testWorkflowHello(ctx Context) (string, error) {
   481  	ao := ActivityOptions{
   482  		ScheduleToStartTimeout: time.Minute,
   483  		StartToCloseTimeout:    time.Minute,
   484  		HeartbeatTimeout:       20 * time.Second,
   485  	}
   486  	ctx = WithActivityOptions(ctx, ao)
   487  
   488  	var result string
   489  	err := ExecuteActivity(ctx, testActivityHello, "world").Get(ctx, &result)
   490  	if err != nil {
   491  		return "", err
   492  	}
   493  	return result, nil
   494  }
   495  
   496  func testWorkflowContext(ctx Context) (string, error) {
   497  	value := ctx.Value(contextKey(testHeader))
   498  	if val, ok := value.(string); ok {
   499  		return val, nil
   500  	}
   501  	return "", fmt.Errorf("context did not propagate to workflow")
   502  }
   503  
   504  func testActivityHello(_ context.Context, msg string) (string, error) {
   505  	return "hello" + "_" + msg, nil
   506  }
   507  
   508  func testActivityContext(ctx context.Context) (string, error) {
   509  	value := ctx.Value(contextKey(testHeader))
   510  	if val, ok := value.(string); ok {
   511  		return val, nil
   512  	}
   513  	return "", fmt.Errorf("context did not propagate to workflow")
   514  }
   515  
   516  func testActivityCanceled(ctx context.Context) (int32, error) {
   517  	info := GetActivityInfo(ctx)
   518  	if info.Attempt < 2 {
   519  		return int32(-1), NewCanceledError("details")
   520  	}
   521  	return info.Attempt, nil
   522  }
   523  
   524  func testWorkflowHeartbeat(ctx Context, msg string, waitTime time.Duration) (string, error) {
   525  	ao := ActivityOptions{
   526  		ScheduleToStartTimeout: time.Minute,
   527  		StartToCloseTimeout:    time.Minute,
   528  		HeartbeatTimeout:       20 * time.Second,
   529  	}
   530  	ctx = WithActivityOptions(ctx, ao)
   531  
   532  	var result string
   533  	err := ExecuteActivity(ctx, testActivityHeartbeat, msg, waitTime).Get(ctx, &result)
   534  	if err != nil {
   535  		return "", err
   536  	}
   537  	return result, nil
   538  }
   539  
   540  func testActivityHeartbeat(ctx context.Context, msg string, waitTime time.Duration) (string, error) {
   541  	GetActivityLogger(ctx).Info("testActivityHeartbeat start",
   542  		zap.String("msg", msg), zap.Duration("waitTime", waitTime))
   543  
   544  	currWaitTime := time.Duration(0)
   545  	for currWaitTime < waitTime {
   546  		RecordActivityHeartbeat(ctx)
   547  		select {
   548  		case <-ctx.Done():
   549  			// We have been cancelled.
   550  			return "", ctx.Err()
   551  		default:
   552  			// We are not cancelled yet.
   553  		}
   554  
   555  		sleepDuration := time.Second
   556  		if currWaitTime+sleepDuration > waitTime {
   557  			sleepDuration = waitTime - currWaitTime
   558  		}
   559  		time.Sleep(sleepDuration)
   560  		currWaitTime += sleepDuration
   561  	}
   562  
   563  	return "heartbeat_" + msg, nil
   564  }
   565  
   566  func (s *WorkflowTestSuiteUnitTest) Test_SideEffect() {
   567  	value := 12345
   568  	workflowFn := func(ctx Context) error {
   569  		ctx = WithActivityOptions(ctx, s.activityOptions)
   570  		se := SideEffect(ctx, func(ctx Context) interface{} {
   571  			return value
   572  		})
   573  		var v int
   574  		if err := se.Get(&v); err != nil {
   575  			return err
   576  		}
   577  		if v != value {
   578  			return errors.New("unexpected")
   579  		}
   580  		f := ExecuteActivity(ctx, testActivityHello, "msg1")
   581  		err := f.Get(ctx, nil) // wait for result
   582  		return err
   583  	}
   584  
   585  	env := s.NewTestWorkflowEnvironment()
   586  	env.RegisterWorkflow(workflowFn)
   587  	env.RegisterActivity(testActivityHello)
   588  
   589  	env.ExecuteWorkflow(workflowFn)
   590  
   591  	s.True(env.IsWorkflowCompleted())
   592  	s.Nil(env.GetWorkflowError())
   593  }
   594  
   595  func (s *WorkflowTestSuiteUnitTest) Test_ChildWorkflow_Basic() {
   596  	workflowFn := func(ctx Context) (string, error) {
   597  		ctx = WithActivityOptions(ctx, s.activityOptions)
   598  		var helloActivityResult string
   599  		err := ExecuteActivity(ctx, testActivityHello, "activity").Get(ctx, &helloActivityResult)
   600  		if err != nil {
   601  			return "", err
   602  		}
   603  
   604  		cwo := ChildWorkflowOptions{ExecutionStartToCloseTimeout: time.Minute}
   605  		ctx = WithChildWorkflowOptions(ctx, cwo)
   606  		var helloWorkflowResult string
   607  		err = ExecuteChildWorkflow(ctx, testWorkflowHello).Get(ctx, &helloWorkflowResult)
   608  		if err != nil {
   609  			return "", err
   610  		}
   611  
   612  		return helloActivityResult + " " + helloWorkflowResult, nil
   613  	}
   614  
   615  	env := s.NewTestWorkflowEnvironment()
   616  	env.RegisterWorkflow(workflowFn)
   617  	env.RegisterWorkflow(testWorkflowHello)
   618  	env.RegisterActivity(testActivityHello)
   619  	env.ExecuteWorkflow(workflowFn)
   620  
   621  	s.True(env.IsWorkflowCompleted())
   622  	s.NoError(env.GetWorkflowError())
   623  	var actualResult string
   624  	s.NoError(env.GetWorkflowResult(&actualResult))
   625  	s.Equal("hello_activity hello_world", actualResult)
   626  }
   627  
   628  func (s *WorkflowTestSuiteUnitTest) Test_ChildWorkflow_Basic_WithDataConverter() {
   629  	workflowFn := func(ctx Context) (string, error) {
   630  		ctx = WithActivityOptions(ctx, s.activityOptions)
   631  		var helloActivityResult string
   632  		err := ExecuteActivity(ctx, testActivityHello, "activity").Get(ctx, &helloActivityResult)
   633  		if err != nil {
   634  			return "", err
   635  		}
   636  
   637  		cwo := ChildWorkflowOptions{ExecutionStartToCloseTimeout: time.Minute}
   638  		ctx = WithChildWorkflowOptions(ctx, cwo)
   639  		var helloWorkflowResult string
   640  		ctx = WithDataConverter(ctx, newTestDataConverter())
   641  		err = ExecuteChildWorkflow(ctx, testWorkflowHello).Get(ctx, &helloWorkflowResult)
   642  		if err != nil {
   643  			return "", err
   644  		}
   645  
   646  		return helloActivityResult + " " + helloWorkflowResult, nil
   647  	}
   648  
   649  	env := s.NewTestWorkflowEnvironment()
   650  	env.RegisterWorkflow(workflowFn)
   651  	env.RegisterWorkflow(testWorkflowHello)
   652  	env.RegisterActivity(testActivityHello)
   653  
   654  	env.ExecuteWorkflow(workflowFn)
   655  
   656  	s.True(env.IsWorkflowCompleted())
   657  	s.NoError(env.GetWorkflowError())
   658  	var actualResult string
   659  	s.NoError(env.GetWorkflowResult(&actualResult))
   660  	s.Equal("hello_activity hello_world", actualResult)
   661  }
   662  
   663  func (s *WorkflowTestSuiteUnitTest) Test_ChildWorkflowCancel() {
   664  	workflowFn := func(ctx Context) error {
   665  		cwo := ChildWorkflowOptions{
   666  			ExecutionStartToCloseTimeout: time.Minute,
   667  			WaitForCancellation:          true,
   668  		}
   669  		ctx = WithChildWorkflowOptions(ctx, cwo)
   670  		ctx1, cancel1 := WithCancel(ctx)
   671  		ctx2, cancel2 := WithCancel(ctx)
   672  		f1 := ExecuteChildWorkflow(ctx1, testWorkflowHeartbeat, "fast", time.Millisecond)
   673  		f2 := ExecuteChildWorkflow(ctx2, testWorkflowHeartbeat, "slow", time.Hour)
   674  
   675  		NewSelector(ctx).AddFuture(f1, func(f Future) {
   676  			cancel2()
   677  		}).AddFuture(f2, func(f Future) {
   678  			cancel1()
   679  		}).Select(ctx)
   680  
   681  		return nil
   682  	}
   683  
   684  	env := s.NewTestWorkflowEnvironment()
   685  	env.RegisterWorkflow(workflowFn)
   686  	env.RegisterWorkflow(testWorkflowHeartbeat)
   687  	env.RegisterActivity(testActivityHeartbeat)
   688  
   689  	env.ExecuteWorkflow(workflowFn)
   690  
   691  	s.True(env.IsWorkflowCompleted())
   692  	s.Nil(env.GetWorkflowError())
   693  }
   694  
   695  func (s *WorkflowTestSuiteUnitTest) Test_ChildWorkflow_Mock() {
   696  	workflowFn := func(ctx Context) (string, error) {
   697  		ctx = WithActivityOptions(ctx, s.activityOptions)
   698  		var helloActivityResult string
   699  		err := ExecuteActivity(ctx, testActivityHello, "activity").Get(ctx, &helloActivityResult)
   700  		if err != nil {
   701  			return "", err
   702  		}
   703  
   704  		cwo := ChildWorkflowOptions{ExecutionStartToCloseTimeout: time.Minute}
   705  		ctx = WithChildWorkflowOptions(ctx, cwo)
   706  		var helloWorkflowResult string
   707  		err = ExecuteChildWorkflow(ctx, testWorkflowHello).Get(ctx, &helloWorkflowResult)
   708  		if err != nil {
   709  			return "", err
   710  		}
   711  
   712  		cwo = ChildWorkflowOptions{WorkflowID: "workflow-id", ExecutionStartToCloseTimeout: time.Minute}
   713  		ctx = WithChildWorkflowOptions(ctx, cwo)
   714  		var heartbeatWorkflowResult string
   715  		err = ExecuteChildWorkflow(ctx, testWorkflowHeartbeat, "slow", time.Hour).Get(ctx, &heartbeatWorkflowResult)
   716  		if err != nil {
   717  			return "", err
   718  		}
   719  
   720  		return helloActivityResult + " " + helloWorkflowResult + " " + heartbeatWorkflowResult, nil
   721  	}
   722  
   723  	env := s.NewTestWorkflowEnvironment()
   724  	env.RegisterWorkflow(workflowFn)
   725  	env.RegisterWorkflow(testWorkflowHello)
   726  	env.RegisterWorkflow(testWorkflowHeartbeat)
   727  	env.RegisterActivity(testActivityHeartbeat)
   728  	env.RegisterActivity(testActivityHello)
   729  
   730  	env.OnActivity(testActivityHello, mock.Anything, mock.Anything).Return("mock_msg", nil)
   731  	env.OnWorkflow(testWorkflowHeartbeat, mock.Anything, mock.Anything, mock.Anything).
   732  		Run(func(a mock.Arguments) {
   733  			ctx := a.Get(0).(Context)
   734  
   735  			// Ensure GetWorkflowInfo is usable in mocks.
   736  			info := GetWorkflowInfo(ctx)
   737  			s.Equal("workflow-id", info.WorkflowExecution.ID)
   738  		}).
   739  		Return("mock_heartbeat", nil)
   740  	env.ExecuteWorkflow(workflowFn)
   741  
   742  	s.True(env.IsWorkflowCompleted())
   743  	s.NoError(env.GetWorkflowError())
   744  	var actualResult string
   745  	s.NoError(env.GetWorkflowResult(&actualResult))
   746  	s.Equal("mock_msg mock_msg mock_heartbeat", actualResult)
   747  }
   748  
   749  // Test_ChildWorkflow_Mock_Panic_GetChildWorkflowExecution verifies that
   750  // ExecuteChildWorkflow(...).GetChildWorkflowExecution().Get() doesn't block forever when mock panics
   751  func (s *WorkflowTestSuiteUnitTest) Test_ChildWorkflow_Mock_Panic_GetChildWorkflowExecution() {
   752  	workflowFn := func(ctx Context) (string, error) {
   753  		cwo := ChildWorkflowOptions{ExecutionStartToCloseTimeout: time.Minute}
   754  		ctx = WithChildWorkflowOptions(ctx, cwo)
   755  		var helloWorkflowResult string
   756  		childWorkflow := ExecuteChildWorkflow(ctx, testWorkflowHello)
   757  		childWorkflow.GetChildWorkflowExecution().Get(ctx, nil)
   758  		err := childWorkflow.Get(ctx, &helloWorkflowResult)
   759  		if err != nil {
   760  			return "", err
   761  		}
   762  		return helloWorkflowResult, nil
   763  	}
   764  
   765  	env := s.NewTestWorkflowEnvironment()
   766  	env.RegisterWorkflowWithOptions(testWorkflowHello, RegisterWorkflowOptions{Name: "testWorkflowHello"})
   767  	env.RegisterWorkflow(workflowFn)
   768  	env.OnWorkflow(testWorkflowHello, mock.Anything, mock.Anything, mock.Anything).
   769  		Return("mock_result", nil, "extra_argument") // extra arg causes panic
   770  	env.ExecuteWorkflow(workflowFn)
   771  
   772  	s.True(env.IsWorkflowCompleted())
   773  	workflowError := env.GetWorkflowError()
   774  	s.Error(workflowError)
   775  	s.Equal("mock of testWorkflowHello has incorrect number of returns, expected 2, but actual is 3",
   776  		workflowError.Error())
   777  }
   778  
   779  func (s *WorkflowTestSuiteUnitTest) Test_ChildWorkflow_StartFailed() {
   780  	workflowFn := func(ctx Context) (string, error) {
   781  		cwo := ChildWorkflowOptions{ExecutionStartToCloseTimeout: time.Minute}
   782  		ctx = WithChildWorkflowOptions(ctx, cwo)
   783  		err := ExecuteChildWorkflow(ctx, testWorkflowHello).GetChildWorkflowExecution().Get(ctx, nil)
   784  		if err != nil {
   785  			return "", errors.New("fail to start child")
   786  		}
   787  
   788  		return "should-not-go-here", nil
   789  	}
   790  
   791  	env := s.NewTestWorkflowEnvironment()
   792  	env.RegisterWorkflow(workflowFn)
   793  	env.RegisterWorkflow(testWorkflowHello)
   794  	env.OnWorkflow(testWorkflowHello, mock.Anything).Return("", ErrMockStartChildWorkflowFailed)
   795  	env.ExecuteWorkflow(workflowFn)
   796  
   797  	s.True(env.IsWorkflowCompleted())
   798  	s.Error(env.GetWorkflowError())
   799  	s.Equal("fail to start child", env.GetWorkflowError().Error())
   800  }
   801  
   802  func (s *WorkflowTestSuiteUnitTest) Test_ChildWorkflow_Listener() {
   803  	workflowFn := func(ctx Context) (string, error) {
   804  		ctx = WithActivityOptions(ctx, s.activityOptions)
   805  		var helloActivityResult string
   806  		err := ExecuteActivity(ctx, testActivityHello, "activity").Get(ctx, &helloActivityResult)
   807  		if err != nil {
   808  			return "", err
   809  		}
   810  
   811  		cwo := ChildWorkflowOptions{ExecutionStartToCloseTimeout: time.Minute}
   812  		ctx = WithChildWorkflowOptions(ctx, cwo)
   813  		var helloWorkflowResult string
   814  		err = ExecuteChildWorkflow(ctx, testWorkflowHello).Get(ctx, &helloWorkflowResult)
   815  		if err != nil {
   816  			return "", err
   817  		}
   818  
   819  		return helloActivityResult + " " + helloWorkflowResult, nil
   820  	}
   821  
   822  	env := s.NewTestWorkflowEnvironment()
   823  	env.RegisterWorkflow(workflowFn)
   824  	env.RegisterWorkflowWithOptions(testWorkflowHello, RegisterWorkflowOptions{Name: "testWorkflowHello"})
   825  	env.RegisterActivity(testActivityHello)
   826  
   827  	var childWorkflowName, childWorkflowResult string
   828  	env.SetOnChildWorkflowStartedListener(func(workflowInfo *WorkflowInfo, ctx Context, args Values) {
   829  		childWorkflowName = workflowInfo.WorkflowType.Name
   830  	})
   831  	env.SetOnChildWorkflowCompletedListener(func(workflowInfo *WorkflowInfo, result Value, err error) {
   832  		s.NoError(err)
   833  		s.NoError(result.Get(&childWorkflowResult))
   834  	})
   835  	env.ExecuteWorkflow(workflowFn)
   836  
   837  	s.True(env.IsWorkflowCompleted())
   838  	s.NoError(env.GetWorkflowError())
   839  	var actualResult string
   840  	s.NoError(env.GetWorkflowResult(&actualResult))
   841  	s.Equal("hello_activity hello_world", actualResult)
   842  	s.Equal("hello_world", childWorkflowResult)
   843  	s.Equal("testWorkflowHello", childWorkflowName)
   844  }
   845  
   846  func (s *WorkflowTestSuiteUnitTest) Test_ChildWorkflow_Clock() {
   847  	expected := []string{
   848  		"child: activity completed",
   849  		"parent: 1m timer fired",
   850  		"parent: 10m timer fired",
   851  		"child: 1h timer fired",
   852  		"parent: child completed",
   853  	}
   854  
   855  	var history []string
   856  	mutex := sync.Mutex{}
   857  	addHistory := func(event string) {
   858  		mutex.Lock()
   859  		history = append(history, event)
   860  		mutex.Unlock()
   861  	}
   862  	childWorkflowFn := func(ctx Context) error {
   863  		t1 := NewTimer(ctx, time.Hour)
   864  		ctx = WithActivityOptions(ctx, s.activityOptions)
   865  		f1 := ExecuteActivity(ctx, testActivityHello, "from child workflow")
   866  
   867  		selector := NewSelector(ctx)
   868  		selector.AddFuture(t1, func(f Future) {
   869  			addHistory("child: 1h timer fired")
   870  		}).AddFuture(f1, func(f Future) {
   871  			addHistory("child: activity completed")
   872  		})
   873  
   874  		selector.Select(ctx)
   875  		selector.Select(ctx)
   876  
   877  		t1.Get(ctx, nil)
   878  		f1.Get(ctx, nil)
   879  
   880  		return nil
   881  	}
   882  
   883  	workflowFn := func(ctx Context) error {
   884  		t1 := NewTimer(ctx, time.Minute)
   885  		t2 := NewTimer(ctx, time.Minute*10)
   886  
   887  		cwo := ChildWorkflowOptions{ExecutionStartToCloseTimeout: time.Hour * 2}
   888  		ctx = WithChildWorkflowOptions(ctx, cwo)
   889  		f1 := ExecuteChildWorkflow(ctx, childWorkflowFn)
   890  
   891  		selector := NewSelector(ctx)
   892  		selector.AddFuture(f1, func(f Future) {
   893  			addHistory("parent: child completed")
   894  		}).AddFuture(t1, func(f Future) {
   895  			addHistory("parent: 1m timer fired")
   896  		}).AddFuture(t2, func(f Future) {
   897  			addHistory("parent: 10m timer fired")
   898  		})
   899  
   900  		selector.Select(ctx)
   901  		selector.Select(ctx)
   902  		selector.Select(ctx)
   903  
   904  		return nil
   905  	}
   906  
   907  	env := s.NewTestWorkflowEnvironment()
   908  	env.RegisterWorkflow(workflowFn)
   909  	env.RegisterWorkflow(childWorkflowFn)
   910  	env.RegisterActivity(testActivityHello)
   911  
   912  	env.ExecuteWorkflow(workflowFn)
   913  
   914  	s.True(env.IsWorkflowCompleted())
   915  	s.NoError(env.GetWorkflowError())
   916  	s.Equal(expected, history)
   917  }
   918  
   919  func (s *WorkflowTestSuiteUnitTest) Test_MockActivityWait() {
   920  	workflowFn := func(ctx Context) error {
   921  		t1 := NewTimer(ctx, time.Hour)
   922  		ctx = WithActivityOptions(ctx, s.activityOptions)
   923  		f1 := ExecuteActivity(ctx, testActivityHello, "mock_delay")
   924  
   925  		NewSelector(ctx).AddFuture(t1, func(f Future) {
   926  			// timer fired
   927  		}).AddFuture(f1, func(f Future) {
   928  			// activity completed
   929  		}).Select(ctx)
   930  
   931  		// either t1 or f1 is ready.
   932  		if f1.IsReady() {
   933  			return nil
   934  		}
   935  
   936  		// activity takes too long
   937  		return errors.New("activity takes too long")
   938  	}
   939  
   940  	// no delay to the mock call, workflow should return no error
   941  	env := s.NewTestWorkflowEnvironment()
   942  	env.RegisterWorkflow(workflowFn)
   943  	env.RegisterActivity(testActivityHello)
   944  	env.OnActivity(testActivityHello, mock.Anything, mock.Anything).Return("hello_mock_delayed", nil).Once()
   945  
   946  	env.ExecuteWorkflow(workflowFn)
   947  	s.True(env.IsWorkflowCompleted())
   948  	s.NoError(env.GetWorkflowError())
   949  	env.AssertExpectations(s.T())
   950  
   951  	// delay 10 minutes, which is shorter than the 1 hour timer, so workflow should return no error.
   952  	env = s.NewTestWorkflowEnvironment()
   953  	env.RegisterActivity(testActivityHello)
   954  	env.OnActivity(testActivityHello, mock.Anything, mock.Anything).After(time.Minute*10).Return("hello_mock_delayed", nil).Once()
   955  	env.ExecuteWorkflow(workflowFn)
   956  	s.True(env.IsWorkflowCompleted())
   957  	s.NoError(env.GetWorkflowError())
   958  	env.AssertExpectations(s.T())
   959  
   960  	// delay 2 hours, which is longer than the 1 hour timer, and workflow should return error.
   961  	env = s.NewTestWorkflowEnvironment()
   962  	env.RegisterActivity(testActivityHello)
   963  	env.OnActivity(testActivityHello, mock.Anything, mock.Anything).After(time.Hour*2).Return("hello_mock_delayed", nil).Once()
   964  	env.ExecuteWorkflow(workflowFn)
   965  	s.True(env.IsWorkflowCompleted())
   966  	s.Error(env.GetWorkflowError())
   967  	env.AssertExpectations(s.T())
   968  
   969  	// no mock
   970  	env = s.NewTestWorkflowEnvironment()
   971  	env.RegisterActivity(testActivityHello)
   972  	env.ExecuteWorkflow(workflowFn)
   973  	s.True(env.IsWorkflowCompleted())
   974  	s.NoError(env.GetWorkflowError())
   975  }
   976  
   977  func (s *WorkflowTestSuiteUnitTest) Test_MockActivityWaitFn() {
   978  	workflowFn := func(ctx Context) ([]string, error) {
   979  		ctx = WithActivityOptions(ctx, s.activityOptions)
   980  		var first, second string
   981  		ExecuteActivity(ctx, testActivityHello, "mock_delay_1").Get(ctx, &first)
   982  		ExecuteActivity(ctx, testActivityHello, "mock_delay_2").Get(ctx, &second)
   983  		return []string{first, second}, nil
   984  	}
   985  
   986  	// extract results from the env and compare to expected values
   987  	expectResult := func(env *TestWorkflowEnvironment, expected []string) {
   988  		var result []string
   989  		err := env.GetWorkflowResult(&result)
   990  		s.NoError(err)
   991  		s.Equal(expected, result)
   992  	}
   993  
   994  	// wrap around ExecuteWorkflow call to track env execution time
   995  	expectDuration := func(env *TestWorkflowEnvironment, seconds int, fn func()) {
   996  		before := env.Now()
   997  
   998  		fn()
   999  
  1000  		after := env.Now()
  1001  		expected := time.Second * time.Duration(seconds)
  1002  		s.Truef(after.Sub(before).Round(time.Second) == expected,
  1003  			"Expected %v to be %v after %v, real diff: %v", after, expected, before, after.Sub(before))
  1004  	}
  1005  
  1006  	// multiple different mocked delays and values should work
  1007  	env := s.NewTestWorkflowEnvironment()
  1008  	env.RegisterWorkflow(workflowFn)
  1009  	env.RegisterActivity(testActivityHello)
  1010  	env.OnActivity(testActivityHello, mock.Anything, "mock_delay_1").After(time.Second).Return("one", nil).Once()
  1011  	env.OnActivity(testActivityHello, mock.Anything, "mock_delay_2").After(time.Minute).Return("two", nil).Once()
  1012  	expectDuration(env, 61, func() {
  1013  		env.ExecuteWorkflow(workflowFn)
  1014  	})
  1015  	s.True(env.IsWorkflowCompleted())
  1016  	s.NoError(env.GetWorkflowError())
  1017  	env.AssertExpectations(s.T())
  1018  	expectResult(env, []string{"one", "two"})
  1019  
  1020  	// multiple different dynamically mocked delays and values should work
  1021  	env = s.NewTestWorkflowEnvironment()
  1022  	env.RegisterWorkflow(workflowFn)
  1023  	env.RegisterActivity(testActivityHello)
  1024  
  1025  	afterCount, returnCount := 0, 0
  1026  	afters := []time.Duration{time.Second, time.Minute}
  1027  	returns := []string{"first", "second"}
  1028  	env.OnActivity(testActivityHello, mock.Anything, mock.Anything).
  1029  		AfterFn(func() time.Duration {
  1030  			defer func() { afterCount++ }()
  1031  			return afters[afterCount]
  1032  		}).
  1033  		Return(func(ctx context.Context, msg string) (string, error) {
  1034  			defer func() { returnCount++ }()
  1035  			return returns[returnCount], nil
  1036  		}).Twice()
  1037  	expectDuration(env, 61, func() {
  1038  		env.ExecuteWorkflow(workflowFn)
  1039  	})
  1040  	s.True(env.IsWorkflowCompleted())
  1041  	s.NoError(env.GetWorkflowError())
  1042  	env.AssertExpectations(s.T())
  1043  	expectResult(env, returns)
  1044  }
  1045  
  1046  func (s *WorkflowTestSuiteUnitTest) Test_MockWorkflowWait() {
  1047  	workflowFn := func(ctx Context) error {
  1048  		t1 := NewTimer(ctx, time.Hour)
  1049  		cwo := ChildWorkflowOptions{ExecutionStartToCloseTimeout: time.Hour /* this is currently ignored by test suite */}
  1050  		ctx = WithChildWorkflowOptions(ctx, cwo)
  1051  		f1 := ExecuteChildWorkflow(ctx, testWorkflowHello)
  1052  
  1053  		NewSelector(ctx).AddFuture(t1, func(f Future) {
  1054  			// timer fired
  1055  		}).AddFuture(f1, func(f Future) {
  1056  			// child workflow completed
  1057  		}).Select(ctx)
  1058  
  1059  		// either t1 or f1 is ready.
  1060  		if f1.IsReady() {
  1061  			return nil
  1062  		}
  1063  
  1064  		// child workflow takes too long
  1065  		return errors.New("child workflow takes too long")
  1066  	}
  1067  
  1068  	// no delay to the mock call, workflow should return no error
  1069  	env := s.NewTestWorkflowEnvironment()
  1070  	env.RegisterWorkflow(workflowFn)
  1071  	env.RegisterWorkflow(testWorkflowHello)
  1072  	env.OnWorkflow(testWorkflowHello, mock.Anything).Return("hello_mock_delayed", nil).Once()
  1073  	env.RegisterActivity(testActivityHello)
  1074  	env.ExecuteWorkflow(workflowFn)
  1075  	s.True(env.IsWorkflowCompleted())
  1076  	s.NoError(env.GetWorkflowError())
  1077  	env.AssertExpectations(s.T())
  1078  
  1079  	// delay 10 minutes, which is shorter than the 1 hour timer, so workflow should return no error.
  1080  	env = s.NewTestWorkflowEnvironment()
  1081  	env.RegisterWorkflow(workflowFn)
  1082  	env.RegisterWorkflow(testWorkflowHello)
  1083  	env.OnWorkflow(testWorkflowHello, mock.Anything).After(time.Minute*10).Return("hello_mock_delayed", nil).Once()
  1084  	env.RegisterActivity(testActivityHello)
  1085  	env.ExecuteWorkflow(workflowFn)
  1086  	s.True(env.IsWorkflowCompleted())
  1087  	s.NoError(env.GetWorkflowError())
  1088  	env.AssertExpectations(s.T())
  1089  
  1090  	// delay 2 hours, which is longer than the 1 hour timer, and workflow should return error.
  1091  	env = s.NewTestWorkflowEnvironment()
  1092  	env.RegisterWorkflow(workflowFn)
  1093  	env.RegisterWorkflow(testWorkflowHello)
  1094  	env.OnWorkflow(testWorkflowHello, mock.Anything).After(time.Hour*2).Return("hello_mock_delayed", nil).Once()
  1095  	env.RegisterActivity(testActivityHello)
  1096  	env.ExecuteWorkflow(workflowFn)
  1097  	s.True(env.IsWorkflowCompleted())
  1098  	s.Error(env.GetWorkflowError())
  1099  	env.AssertExpectations(s.T())
  1100  
  1101  	// no mock
  1102  	env = s.NewTestWorkflowEnvironment()
  1103  	env.RegisterWorkflow(workflowFn)
  1104  	env.RegisterWorkflow(testWorkflowHello)
  1105  	env.RegisterActivity(testActivityHello)
  1106  	env.ExecuteWorkflow(workflowFn)
  1107  	s.True(env.IsWorkflowCompleted())
  1108  	s.NoError(env.GetWorkflowError())
  1109  }
  1110  
  1111  func (s *WorkflowTestSuiteUnitTest) Test_MockPanic() {
  1112  	// mock panic, verify that the panic won't be swallowed by our panic handler to detect unexpected mock call.
  1113  	oldLogger := s.GetLogger()
  1114  	s.SetLogger(zap.NewNop()) // use no-op logger to avoid noisy logging by panic
  1115  	env := s.NewTestWorkflowEnvironment()
  1116  	env.OnActivity(testActivityHello, mock.Anything, mock.Anything).
  1117  		Return("hello_mock_panic", nil).
  1118  		Run(func(args mock.Arguments) {
  1119  			panic("mock-panic")
  1120  		})
  1121  	env.RegisterWorkflow(testWorkflowHello)
  1122  	env.RegisterActivity(testActivityHello)
  1123  	env.ExecuteWorkflow(testWorkflowHello)
  1124  	s.True(env.IsWorkflowCompleted())
  1125  	err := env.GetWorkflowError()
  1126  	s.Error(err)
  1127  	s.Contains(err.Error(), "mock-panic")
  1128  	env.AssertExpectations(s.T())
  1129  	s.SetLogger(oldLogger) // restore original logger
  1130  }
  1131  
  1132  func (s *WorkflowTestSuiteUnitTest) Test_ChildWithChild() {
  1133  	env := s.NewTestWorkflowEnvironment()
  1134  	childWorkflowFn := func(ctx Context) error {
  1135  		t1 := NewTimer(ctx, time.Hour)
  1136  		cwo := ChildWorkflowOptions{ExecutionStartToCloseTimeout: time.Hour /* this is currently ignored by test suite */}
  1137  		ctx = WithChildWorkflowOptions(ctx, cwo)
  1138  		f1 := ExecuteChildWorkflow(ctx, testWorkflowHello)
  1139  
  1140  		NewSelector(ctx).AddFuture(t1, func(f Future) {
  1141  			// timer fired
  1142  		}).AddFuture(f1, func(f Future) {
  1143  			// child workflow completed
  1144  		}).Select(ctx)
  1145  
  1146  		// either t1 or f1 is ready.
  1147  		if f1.IsReady() {
  1148  			return nil
  1149  		}
  1150  
  1151  		// child workflow takes too long
  1152  		return errors.New("child workflow takes too long")
  1153  	}
  1154  
  1155  	workflowFn := func(ctx Context) error {
  1156  		t1 := NewTimer(ctx, time.Hour)
  1157  		cwo := ChildWorkflowOptions{ExecutionStartToCloseTimeout: time.Hour /* this is currently ignored by test suite */}
  1158  		ctx = WithChildWorkflowOptions(ctx, cwo)
  1159  		f1 := ExecuteChildWorkflow(ctx, childWorkflowFn) // execute child workflow which in turn execute another child workflow
  1160  
  1161  		NewSelector(ctx).AddFuture(t1, func(f Future) {
  1162  			// timer fired
  1163  		}).AddFuture(f1, func(f Future) {
  1164  			// child workflow completed
  1165  		}).Select(ctx)
  1166  
  1167  		// either t1 or f1 is ready.
  1168  		if f1.IsReady() {
  1169  			return f1.Get(ctx, nil)
  1170  		}
  1171  
  1172  		// child workflow takes too long
  1173  		return errors.New("child workflow takes too long")
  1174  	}
  1175  
  1176  	env.RegisterWorkflow(childWorkflowFn)
  1177  	env.RegisterWorkflow(workflowFn)
  1178  	// no delay to the mock call, workflow should return no error
  1179  	env.RegisterWorkflow(testWorkflowHello)
  1180  	env.RegisterActivity(testActivityHello)
  1181  	env.OnWorkflow(testWorkflowHello, mock.Anything).Return("hello_mock_delayed", nil).Once()
  1182  	env.ExecuteWorkflow(workflowFn)
  1183  	s.True(env.IsWorkflowCompleted())
  1184  	s.NoError(env.GetWorkflowError())
  1185  	env.AssertExpectations(s.T())
  1186  
  1187  	// delay 10 minutes, which is shorter than the 1 hour timer, so workflow should return no error.
  1188  	env = s.NewTestWorkflowEnvironment()
  1189  	env.RegisterWorkflow(workflowFn)
  1190  	env.RegisterWorkflow(childWorkflowFn)
  1191  
  1192  	// no delay to the mock call, workflow should return no error
  1193  	env.RegisterWorkflow(testWorkflowHello)
  1194  	env.RegisterActivity(testActivityHello)
  1195  
  1196  	env.OnWorkflow(testWorkflowHello, mock.Anything).After(time.Minute*10).Return("hello_mock_delayed", nil).Once()
  1197  	env.ExecuteWorkflow(workflowFn)
  1198  	s.True(env.IsWorkflowCompleted())
  1199  	s.NoError(env.GetWorkflowError())
  1200  	env.AssertExpectations(s.T())
  1201  
  1202  	// delay 2 hours, which is longer than the 1 hour timer, and workflow should return error.
  1203  	env = s.NewTestWorkflowEnvironment()
  1204  	env.RegisterWorkflow(workflowFn)
  1205  	env.RegisterWorkflow(childWorkflowFn)
  1206  	// no delay to the mock call, workflow should return no error
  1207  	env.RegisterWorkflow(testWorkflowHello)
  1208  	env.RegisterActivity(testActivityHello)
  1209  
  1210  	env.OnWorkflow(testWorkflowHello, mock.Anything).After(time.Hour*2).Return("hello_mock_delayed", nil).Once()
  1211  	env.ExecuteWorkflow(workflowFn)
  1212  	s.True(env.IsWorkflowCompleted())
  1213  	s.Error(env.GetWorkflowError())
  1214  	env.AssertExpectations(s.T())
  1215  
  1216  	// no mock
  1217  	env = s.NewTestWorkflowEnvironment()
  1218  	env.RegisterWorkflow(workflowFn)
  1219  	env.RegisterWorkflow(childWorkflowFn)
  1220  	// no delay to the mock call, workflow should return no error
  1221  	env.RegisterWorkflow(testWorkflowHello)
  1222  	env.RegisterActivity(testActivityHello)
  1223  
  1224  	env.ExecuteWorkflow(workflowFn)
  1225  	s.True(env.IsWorkflowCompleted())
  1226  	s.NoError(env.GetWorkflowError())
  1227  }
  1228  
  1229  func (s *WorkflowTestSuiteUnitTest) Test_GetVersion() {
  1230  	oldActivity := func(ctx context.Context, msg string) (string, error) {
  1231  		return "hello" + "_" + msg, nil
  1232  	}
  1233  	newActivity := func(ctx context.Context, msg string) (string, error) {
  1234  		return "hello" + "_" + msg, nil
  1235  	}
  1236  	workflowFn := func(ctx Context) error {
  1237  		ctx = WithActivityOptions(ctx, s.activityOptions)
  1238  		var f Future
  1239  		v := GetVersion(ctx, "test_change_id", DefaultVersion, 2)
  1240  		if v == DefaultVersion {
  1241  			f = ExecuteActivity(ctx, oldActivity, "ols_msg")
  1242  		} else {
  1243  			f = ExecuteActivity(ctx, newActivity, "new_msg")
  1244  		}
  1245  		err := f.Get(ctx, nil) // wait for result
  1246  		if err != nil {
  1247  			return err
  1248  		}
  1249  
  1250  		// test searchable change version
  1251  		wfInfo := GetWorkflowInfo(ctx)
  1252  		s.NotNil(wfInfo.SearchAttributes)
  1253  		changeVersionsBytes, ok := wfInfo.SearchAttributes.IndexedFields[CadenceChangeVersion]
  1254  		s.True(ok)
  1255  		var changeVersions []string
  1256  		err = json.Unmarshal(changeVersionsBytes, &changeVersions)
  1257  		s.NoError(err)
  1258  		s.Equal(1, len(changeVersions))
  1259  		s.Equal("test_change_id-2", changeVersions[0])
  1260  
  1261  		return err
  1262  	}
  1263  
  1264  	env := s.NewTestWorkflowEnvironment()
  1265  	env.RegisterWorkflow(workflowFn)
  1266  	env.RegisterActivity(oldActivity)
  1267  	env.RegisterActivity(newActivity)
  1268  	env.OnActivity(newActivity, mock.Anything, "new_msg").Return("hello new_mock_msg", nil).Once()
  1269  	env.ExecuteWorkflow(workflowFn)
  1270  
  1271  	s.True(env.IsWorkflowCompleted())
  1272  	s.Nil(env.GetWorkflowError())
  1273  	env.AssertExpectations(s.T())
  1274  }
  1275  
  1276  func (s *WorkflowTestSuiteUnitTest) Test_MockGetVersion() {
  1277  	oldActivity := func(ctx context.Context, msg string) (string, error) {
  1278  		return "hello" + "_" + msg, nil
  1279  	}
  1280  	newActivity := func(ctx context.Context, msg string) (string, error) {
  1281  		return "hello" + "_" + msg, nil
  1282  	}
  1283  	workflowFn := func(ctx Context) (string, error) {
  1284  		ctx = WithActivityOptions(ctx, s.activityOptions)
  1285  		var f Future
  1286  		v1 := GetVersion(ctx, "change_1", DefaultVersion, 2)
  1287  		if v1 == DefaultVersion {
  1288  			f = ExecuteActivity(ctx, oldActivity, "old1")
  1289  		} else {
  1290  			f = ExecuteActivity(ctx, newActivity, "new1")
  1291  		}
  1292  		var ret1 string
  1293  		err := f.Get(ctx, &ret1) // wait for result
  1294  		if err != nil {
  1295  			return "", err
  1296  		}
  1297  
  1298  		v2 := GetVersion(ctx, "change_2", DefaultVersion, 2)
  1299  		if v2 == DefaultVersion {
  1300  			f = ExecuteActivity(ctx, oldActivity, "old2")
  1301  		} else {
  1302  			f = ExecuteActivity(ctx, newActivity, "new2")
  1303  		}
  1304  		var ret2 string
  1305  		err = f.Get(ctx, &ret2) // wait for result
  1306  		if err != nil {
  1307  			return "", err
  1308  		}
  1309  
  1310  		// test searchable change version
  1311  		wfInfo := GetWorkflowInfo(ctx)
  1312  		s.NotNil(wfInfo.SearchAttributes)
  1313  		changeVersionsBytes, ok := wfInfo.SearchAttributes.IndexedFields[CadenceChangeVersion]
  1314  		s.True(ok)
  1315  		var changeVersions []string
  1316  		err = json.Unmarshal(changeVersionsBytes, &changeVersions)
  1317  		s.NoError(err)
  1318  		s.Equal(2, len(changeVersions))
  1319  		s.Equal("change_2-2", changeVersions[0])
  1320  		s.Equal("change_1--1", changeVersions[1])
  1321  
  1322  		return ret1 + ret2, err
  1323  	}
  1324  
  1325  	env := s.NewTestWorkflowEnvironment()
  1326  	env.RegisterWorkflow(workflowFn)
  1327  	env.RegisterActivity(oldActivity)
  1328  	env.RegisterActivity(newActivity)
  1329  
  1330  	env.OnGetVersion("change_1", DefaultVersion, 2).Return(func(string, Version, Version) Version {
  1331  		return DefaultVersion
  1332  	})
  1333  	env.OnGetVersion(mock.Anything, DefaultVersion, 2).Return(Version(2))
  1334  	env.ExecuteWorkflow(workflowFn)
  1335  
  1336  	s.True(env.IsWorkflowCompleted())
  1337  	s.Nil(env.GetWorkflowError())
  1338  	var ret string
  1339  	s.NoError(env.GetWorkflowResult(&ret))
  1340  	s.Equal("hello_old1hello_new2", ret)
  1341  	env.AssertExpectations(s.T())
  1342  }
  1343  
  1344  func (s *WorkflowTestSuiteUnitTest) Test_UpsertSearchAttributes_ReservedKey() {
  1345  	workflowFn := func(ctx Context) error {
  1346  		attr := map[string]interface{}{
  1347  			CadenceChangeVersion: "some change version",
  1348  		}
  1349  		err := UpsertSearchAttributes(ctx, attr)
  1350  		s.Error(err)
  1351  
  1352  		wfInfo := GetWorkflowInfo(ctx)
  1353  		s.Nil(wfInfo.SearchAttributes)
  1354  		return nil
  1355  	}
  1356  	env := s.NewTestWorkflowEnvironment()
  1357  	env.RegisterWorkflow(workflowFn)
  1358  
  1359  	env.ExecuteWorkflow(workflowFn)
  1360  	s.True(env.IsWorkflowCompleted())
  1361  	s.Nil(env.GetWorkflowError())
  1362  	env.AssertExpectations(s.T())
  1363  }
  1364  
  1365  func (s *WorkflowTestSuiteUnitTest) Test_MockUpsertSearchAttributes() {
  1366  	workflowFn := func(ctx Context) error {
  1367  		attr := map[string]interface{}{}
  1368  		err := UpsertSearchAttributes(ctx, attr)
  1369  		s.Error(err)
  1370  
  1371  		wfInfo := GetWorkflowInfo(ctx)
  1372  		s.Nil(wfInfo.SearchAttributes)
  1373  
  1374  		attr["CustomIntField"] = 1
  1375  		err = UpsertSearchAttributes(ctx, attr)
  1376  		s.NoError(err)
  1377  
  1378  		wfInfo = GetWorkflowInfo(ctx)
  1379  		s.NotNil(wfInfo.SearchAttributes)
  1380  		valBytes := wfInfo.SearchAttributes.IndexedFields["CustomIntField"]
  1381  		var result int
  1382  		NewValue(valBytes).Get(&result)
  1383  		s.Equal(1, result)
  1384  
  1385  		return nil
  1386  	}
  1387  
  1388  	// no mock
  1389  	env := s.NewTestWorkflowEnvironment()
  1390  	env.RegisterWorkflow(workflowFn)
  1391  
  1392  	env.ExecuteWorkflow(workflowFn)
  1393  	s.True(env.IsWorkflowCompleted())
  1394  	s.Nil(env.GetWorkflowError())
  1395  	env.AssertExpectations(s.T())
  1396  
  1397  	// has mock
  1398  	env = s.NewTestWorkflowEnvironment()
  1399  	env.OnUpsertSearchAttributes(map[string]interface{}{}).Return(errors.New("empty")).Once()
  1400  	env.OnUpsertSearchAttributes(map[string]interface{}{"CustomIntField": 1}).Return(nil).Once()
  1401  
  1402  	env.ExecuteWorkflow(workflowFn)
  1403  	s.True(env.IsWorkflowCompleted())
  1404  	s.Nil(env.GetWorkflowError())
  1405  	env.AssertExpectations(s.T())
  1406  
  1407  	// mock no return
  1408  	env = s.NewTestWorkflowEnvironment()
  1409  	env.OnUpsertSearchAttributes(map[string]interface{}{}).Once()
  1410  	env.OnUpsertSearchAttributes(map[string]interface{}{"CustomIntField": 1}).Once()
  1411  
  1412  	env.ExecuteWorkflow(workflowFn)
  1413  	s.True(env.IsWorkflowCompleted())
  1414  	s.Nil(env.GetWorkflowError())
  1415  	env.AssertExpectations(s.T())
  1416  
  1417  	// mix no-mock and mock is not support
  1418  }
  1419  
  1420  func (s *WorkflowTestSuiteUnitTest) Test_MockUpsertSearchAttributes_OnError() {
  1421  	workflowFn := func(ctx Context) error {
  1422  		attr := map[string]interface{}{}
  1423  		attr["CustomIntField"] = 1
  1424  		err := UpsertSearchAttributes(ctx, attr)
  1425  		s.Error(err)
  1426  		return err
  1427  	}
  1428  
  1429  	env := s.NewTestWorkflowEnvironment()
  1430  	env.OnUpsertSearchAttributes(map[string]interface{}{"CustomIntField": 1}).Return(errors.New("test error")).Once()
  1431  
  1432  	env.ExecuteWorkflow(workflowFn)
  1433  	s.True(env.IsWorkflowCompleted())
  1434  	s.Error(env.GetWorkflowError())
  1435  	env.AssertExpectations(s.T())
  1436  }
  1437  
  1438  func (s *WorkflowTestSuiteUnitTest) Test_ActivityWithThriftTypes() {
  1439  	actualValues := []string{}
  1440  	retVal := &shared.WorkflowExecution{WorkflowId: common.StringPtr("retwID2"), RunId: common.StringPtr("retrID2")}
  1441  
  1442  	// Passing one argument
  1443  	activitySingleFn := func(ctx context.Context, wf *shared.WorkflowExecution) (*shared.WorkflowExecution, error) {
  1444  		actualValues = append(actualValues, wf.GetWorkflowId())
  1445  		actualValues = append(actualValues, wf.GetRunId())
  1446  		return retVal, nil
  1447  	}
  1448  
  1449  	input := &shared.WorkflowExecution{WorkflowId: common.StringPtr("wID1"), RunId: common.StringPtr("rID1")}
  1450  	env := s.NewTestActivityEnvironment()
  1451  	env.RegisterActivity(activitySingleFn)
  1452  	blob, err := env.ExecuteActivity(activitySingleFn, input)
  1453  	s.NoError(err)
  1454  	var ret *shared.WorkflowExecution
  1455  	blob.Get(&ret)
  1456  	s.Equal(retVal, ret)
  1457  
  1458  	// Passing more than one argument
  1459  	activityDoubleArgFn := func(ctx context.Context, wf *shared.WorkflowExecution, t *shared.WorkflowType) (*shared.WorkflowExecution, error) {
  1460  		actualValues = append(actualValues, wf.GetWorkflowId())
  1461  		actualValues = append(actualValues, wf.GetRunId())
  1462  		actualValues = append(actualValues, t.GetName())
  1463  		return retVal, nil
  1464  	}
  1465  
  1466  	input = &shared.WorkflowExecution{WorkflowId: common.StringPtr("wID2"), RunId: common.StringPtr("rID3")}
  1467  	wt := &shared.WorkflowType{Name: common.StringPtr("wType")}
  1468  	env = s.NewTestActivityEnvironment()
  1469  	env.RegisterActivity(activityDoubleArgFn)
  1470  	blob, err = env.ExecuteActivity(activityDoubleArgFn, input, wt)
  1471  	s.NoError(err)
  1472  	blob.Get(&ret)
  1473  	s.Equal(retVal, ret)
  1474  
  1475  	expectedValues := []string{
  1476  		"wID1",
  1477  		"rID1",
  1478  		"wID2",
  1479  		"rID3",
  1480  		"wType",
  1481  	}
  1482  	s.EqualValues(expectedValues, actualValues)
  1483  }
  1484  
  1485  func (s *WorkflowTestSuiteUnitTest) Test_ActivityRegistration() {
  1486  	activityFn := func(msg string) (string, error) {
  1487  		return msg, nil
  1488  	}
  1489  	activityAlias := "some-random-activity-alias"
  1490  
  1491  	env := s.NewTestActivityEnvironment()
  1492  	env.RegisterActivityWithOptions(activityFn, RegisterActivityOptions{Name: activityAlias})
  1493  	input := "some random input"
  1494  
  1495  	encodedValue, err := env.ExecuteActivity(activityFn, input)
  1496  	s.NoError(err)
  1497  	output := ""
  1498  	encodedValue.Get(&output)
  1499  	s.Equal(input, output)
  1500  
  1501  	encodedValue, err = env.ExecuteActivity(activityAlias, input)
  1502  	s.NoError(err)
  1503  	output = ""
  1504  	encodedValue.Get(&output)
  1505  	s.Equal(input, output)
  1506  }
  1507  
  1508  func (s *WorkflowTestSuiteUnitTest) Test_WorkflowRegistration() {
  1509  	workflowFn := func(ctx Context) error {
  1510  		return nil
  1511  	}
  1512  	workflowAlias := "some-random-workflow-alias"
  1513  
  1514  	env := s.NewTestWorkflowEnvironment()
  1515  
  1516  	env.ExecuteWorkflow(workflowFn)
  1517  
  1518  	env = s.NewTestWorkflowEnvironment()
  1519  	env.RegisterWorkflowWithOptions(workflowFn, RegisterWorkflowOptions{Name: workflowAlias})
  1520  
  1521  	env.ExecuteWorkflow(workflowAlias)
  1522  }
  1523  
  1524  func (s *WorkflowTestSuiteUnitTest) Test_ActivityFriendlyName() {
  1525  	activityFn := func(msg string) (string, error) {
  1526  		return "hello_" + msg, nil
  1527  	}
  1528  
  1529  	workflowFn := func(ctx Context) error {
  1530  		ctx = WithActivityOptions(ctx, s.activityOptions)
  1531  		var result string
  1532  		err := ExecuteActivity(ctx, activityFn, "friendly_name").Get(ctx, &result)
  1533  		if err != nil {
  1534  			return err
  1535  		}
  1536  		err = ExecuteActivity(ctx, "foo", "friendly_name").Get(ctx, &result)
  1537  		return err
  1538  	}
  1539  
  1540  	env := s.NewTestWorkflowEnvironment()
  1541  	env.RegisterWorkflow(workflowFn)
  1542  	env.RegisterActivityWithOptions(activityFn, RegisterActivityOptions{Name: "foo"})
  1543  	var called []string
  1544  	env.SetOnActivityStartedListener(func(activityInfo *ActivityInfo, ctx context.Context, args Values) {
  1545  		called = append(called, activityInfo.ActivityType.Name)
  1546  	})
  1547  
  1548  	env.ExecuteWorkflow(workflowFn)
  1549  
  1550  	s.True(env.IsWorkflowCompleted())
  1551  	s.NoError(env.GetWorkflowError())
  1552  	s.Equal(2, len(called))
  1553  	s.Equal("foo", called[0])
  1554  	s.Equal("foo", called[1])
  1555  }
  1556  
  1557  func (s *WorkflowTestSuiteUnitTest) Test_WorkflowFriendlyName() {
  1558  
  1559  	workflowFn := func(ctx Context) error {
  1560  		cwo := ChildWorkflowOptions{ExecutionStartToCloseTimeout: time.Hour /* this is currently ignored by test suite */}
  1561  		ctx = WithChildWorkflowOptions(ctx, cwo)
  1562  		var result string
  1563  		err := ExecuteChildWorkflow(ctx, testWorkflowHello).Get(ctx, &result)
  1564  		if err != nil {
  1565  			return err
  1566  		}
  1567  		err = ExecuteChildWorkflow(ctx, "testWorkflowHello").Get(ctx, &result)
  1568  		return err
  1569  	}
  1570  
  1571  	env := s.NewTestWorkflowEnvironment()
  1572  	env.RegisterWorkflow(workflowFn)
  1573  	env.RegisterWorkflowWithOptions(testWorkflowHello, RegisterWorkflowOptions{Name: "testWorkflowHello"})
  1574  	env.RegisterActivity(testActivityHello)
  1575  	var called []string
  1576  	env.SetOnChildWorkflowStartedListener(func(workflowInfo *WorkflowInfo, ctx Context, args Values) {
  1577  		called = append(called, workflowInfo.WorkflowType.Name)
  1578  	})
  1579  
  1580  	env.ExecuteWorkflow(workflowFn)
  1581  
  1582  	s.True(env.IsWorkflowCompleted())
  1583  	s.NoError(env.GetWorkflowError())
  1584  	s.Equal(2, len(called))
  1585  	s.Equal("testWorkflowHello", called[0])
  1586  	s.Equal("testWorkflowHello", called[1])
  1587  }
  1588  
  1589  func (s *WorkflowTestSuiteUnitTest) Test_WorkflowHeaderContext() {
  1590  
  1591  	workflowFn := func(ctx Context) error {
  1592  		value := ctx.Value(contextKey(testHeader))
  1593  		if val, ok := value.(string); ok {
  1594  			s.Equal("test-data", val)
  1595  		} else {
  1596  			return fmt.Errorf("context did not propagate to workflow")
  1597  		}
  1598  
  1599  		cwo := ChildWorkflowOptions{ExecutionStartToCloseTimeout: time.Hour /* this is currently ignored by test suite */}
  1600  		ctx = WithChildWorkflowOptions(ctx, cwo)
  1601  		var result string
  1602  		if err := ExecuteChildWorkflow(ctx, testWorkflowContext).Get(ctx, &result); err != nil {
  1603  			return err
  1604  		}
  1605  		s.Equal("test-data", result)
  1606  
  1607  		ao := ActivityOptions{
  1608  			ScheduleToStartTimeout: time.Minute,
  1609  			StartToCloseTimeout:    time.Minute,
  1610  			HeartbeatTimeout:       20 * time.Second,
  1611  		}
  1612  		ctx = WithActivityOptions(ctx, ao)
  1613  		if err := ExecuteActivity(ctx, testActivityContext).Get(ctx, &result); err != nil {
  1614  			return err
  1615  		}
  1616  		s.Equal("test-data", result)
  1617  		return nil
  1618  	}
  1619  
  1620  	s.SetContextPropagators([]ContextPropagator{NewStringMapPropagator([]string{testHeader})})
  1621  	s.SetHeader(&shared.Header{
  1622  		Fields: map[string][]byte{
  1623  			testHeader: []byte("test-data"),
  1624  		},
  1625  	})
  1626  
  1627  	env := s.NewTestWorkflowEnvironment()
  1628  	env.RegisterWorkflow(workflowFn)
  1629  	env.ExecuteWorkflow(testWorkflowContext)
  1630  
  1631  	s.True(env.IsWorkflowCompleted())
  1632  	s.NoError(env.GetWorkflowError())
  1633  }
  1634  
  1635  func (s *WorkflowTestSuiteUnitTest) Test_ChildWorkflowContextPropagation() {
  1636  
  1637  	childWorkflowFn := func(ctx Context) error {
  1638  		value := ctx.Value(contextKey(testHeader))
  1639  		if val, ok := value.(string); ok {
  1640  			s.Equal("test-data-for-child", val)
  1641  		} else {
  1642  			return fmt.Errorf("context did not propagate to child workflow")
  1643  		}
  1644  		return nil
  1645  	}
  1646  
  1647  	workflowFn := func(ctx Context) error {
  1648  		value := ctx.Value(contextKey(testHeader))
  1649  		if val, ok := value.(string); ok {
  1650  			s.Equal("test-data", val)
  1651  		} else {
  1652  			return fmt.Errorf("context did not propagate to workflow")
  1653  		}
  1654  
  1655  		cwo := ChildWorkflowOptions{ExecutionStartToCloseTimeout: time.Hour /* this is currently ignored by test suite */}
  1656  		ctx = WithChildWorkflowOptions(ctx, cwo)
  1657  		childCtx := WithValue(ctx, contextKey(testHeader), "test-data-for-child")
  1658  		if err := ExecuteChildWorkflow(childCtx, childWorkflowFn).Get(childCtx, nil); err != nil {
  1659  			return err
  1660  		}
  1661  		return nil
  1662  	}
  1663  
  1664  	s.SetContextPropagators([]ContextPropagator{NewStringMapPropagator([]string{testHeader})})
  1665  	s.SetHeader(&shared.Header{
  1666  		Fields: map[string][]byte{
  1667  			testHeader: []byte("test-data"),
  1668  		},
  1669  	})
  1670  
  1671  	env := s.NewTestWorkflowEnvironment()
  1672  	env.RegisterWorkflow(workflowFn)
  1673  	env.RegisterWorkflow(childWorkflowFn)
  1674  	env.ExecuteWorkflow(workflowFn)
  1675  
  1676  	s.True(env.IsWorkflowCompleted())
  1677  	s.NoError(env.GetWorkflowError())
  1678  }
  1679  
  1680  func (s *WorkflowTestSuiteUnitTest) Test_ActivityFullyQualifiedName() {
  1681  	// TODO (madhu): Add this back once test workflow environment is able to handle panics gracefully
  1682  	// Right now, the panic happens in a different goroutine and there is no way to catch it
  1683  	s.T().Skip()
  1684  	workflowFn := func(ctx Context) error {
  1685  		ctx = WithActivityOptions(ctx, s.activityOptions)
  1686  		var result string
  1687  		fut := ExecuteActivity(ctx, getFunctionName(testActivityHello), "friendly_name")
  1688  		err := fut.Get(ctx, &result)
  1689  		return err
  1690  	}
  1691  
  1692  	env := s.NewTestWorkflowEnvironment()
  1693  	env.RegisterWorkflow(workflowFn)
  1694  	env.ExecuteWorkflow(workflowFn)
  1695  	s.False(env.IsWorkflowCompleted())
  1696  	s.Contains(env.GetWorkflowError().Error(), "Unable to find activityType")
  1697  }
  1698  
  1699  func (s *WorkflowTestSuiteUnitTest) Test_WorkflowUnknownName() {
  1700  	defer func() {
  1701  		if r := recover(); r != nil {
  1702  			s.Contains(r.(error).Error(), "unable to find workflow type")
  1703  		}
  1704  	}()
  1705  	env := s.NewTestWorkflowEnvironment()
  1706  
  1707  	env.ExecuteWorkflow(getFunctionName(testWorkflowHello))
  1708  	s.Fail("Should have panic'ed at ExecuteWorkflow")
  1709  }
  1710  
  1711  func (s *WorkflowTestSuiteUnitTest) Test_QueryWorkflow() {
  1712  	queryType := "state"
  1713  	stateWaitSignal, stateWaitActivity, stateDone := "wait for signal", "wait for activity", "done"
  1714  	workflowFn := func(ctx Context) error {
  1715  		var state string
  1716  		err := SetQueryHandler(ctx, queryType, func(queryInput string) (string, error) {
  1717  			return queryInput + state, nil
  1718  		})
  1719  		if err != nil {
  1720  			return err
  1721  		}
  1722  
  1723  		state = stateWaitSignal
  1724  		var signalData string
  1725  		GetSignalChannel(ctx, "query-signal").Receive(ctx, &signalData)
  1726  
  1727  		state = stateWaitActivity
  1728  		ctx = WithActivityOptions(ctx, s.activityOptions)
  1729  		err = ExecuteActivity(ctx, testActivityHello, "mock_delay").Get(ctx, nil)
  1730  		if err != nil {
  1731  			return err
  1732  		}
  1733  		state = stateDone
  1734  		return err
  1735  	}
  1736  
  1737  	env := s.NewTestWorkflowEnvironment()
  1738  	env.RegisterWorkflow(workflowFn)
  1739  	env.RegisterActivity(testActivityHello)
  1740  
  1741  	verifyStateWithQuery := func(expected string) {
  1742  		encodedValue, err := env.QueryWorkflow(queryType, "input")
  1743  		s.NoError(err)
  1744  		s.NotNil(encodedValue)
  1745  		var state string
  1746  		err = encodedValue.Get(&state)
  1747  		s.NoError(err)
  1748  		s.Equal("input"+expected, state)
  1749  	}
  1750  	env.RegisterDelayedCallback(func() {
  1751  		verifyStateWithQuery(stateWaitSignal)
  1752  		env.SignalWorkflow("query-signal", "hello-query")
  1753  	}, time.Hour)
  1754  	env.OnActivity(testActivityHello, mock.Anything, mock.Anything).After(time.Hour).Return("hello_mock", nil)
  1755  	env.SetOnActivityStartedListener(func(activityInfo *ActivityInfo, ctx context.Context, args Values) {
  1756  		verifyStateWithQuery(stateWaitActivity)
  1757  	})
  1758  	env.ExecuteWorkflow(workflowFn)
  1759  
  1760  	s.True(env.IsWorkflowCompleted())
  1761  	s.NoError(env.GetWorkflowError())
  1762  	env.AssertExpectations(s.T())
  1763  	verifyStateWithQuery(stateDone)
  1764  }
  1765  
  1766  func (s *WorkflowTestSuiteUnitTest) Test_WorkflowWithLocalActivity() {
  1767  	localActivityFn := func(ctx context.Context, name string) (string, error) {
  1768  		return "hello " + name, nil
  1769  	}
  1770  
  1771  	workflowFn := func(ctx Context) (string, error) {
  1772  		ctx = WithLocalActivityOptions(ctx, s.localActivityOptions)
  1773  		var result string
  1774  		f := ExecuteLocalActivity(ctx, localActivityFn, "local_activity")
  1775  		err := f.Get(ctx, &result)
  1776  		return result, err
  1777  	}
  1778  
  1779  	env := s.NewTestWorkflowEnvironment()
  1780  	env.RegisterWorkflow(workflowFn)
  1781  	env.ExecuteWorkflow(workflowFn)
  1782  	s.True(env.IsWorkflowCompleted())
  1783  	s.NoError(env.GetWorkflowError())
  1784  	var result string
  1785  	err := env.GetWorkflowResult(&result)
  1786  	s.NoError(err)
  1787  	s.Equal("hello local_activity", result)
  1788  }
  1789  
  1790  func (s *WorkflowTestSuiteUnitTest) Test_LocalActivity() {
  1791  	localActivityFn := func(ctx context.Context, name string) (string, error) {
  1792  		return "hello " + name, nil
  1793  	}
  1794  
  1795  	env := s.NewTestActivityEnvironment()
  1796  	result, err := env.ExecuteLocalActivity(localActivityFn, "local_activity")
  1797  	s.NoError(err)
  1798  	var laResult string
  1799  	err = result.Get(&laResult)
  1800  	s.NoError(err)
  1801  	s.Equal("hello local_activity", laResult)
  1802  }
  1803  
  1804  func (s *WorkflowTestSuiteUnitTest) Test_WorkflowLocalActivityWithMockAndListeners() {
  1805  	localActivityFn := func(ctx context.Context, name string) (string, error) {
  1806  		return "hello " + name, nil
  1807  	}
  1808  
  1809  	cancelledLocalActivityFn := func() error {
  1810  		time.Sleep(time.Second)
  1811  		return nil
  1812  	}
  1813  
  1814  	workflowFn := func(ctx Context) (string, error) {
  1815  		ctx = WithLocalActivityOptions(ctx, s.localActivityOptions)
  1816  		var result string
  1817  		f := ExecuteLocalActivity(ctx, localActivityFn, "local_activity")
  1818  		ctx2, cancel := WithCancel(ctx)
  1819  		f2 := ExecuteLocalActivity(ctx2, cancelledLocalActivityFn)
  1820  
  1821  		NewSelector(ctx).AddFuture(f, func(f Future) {
  1822  			cancel()
  1823  		}).AddFuture(f2, func(f Future) {
  1824  
  1825  		}).Select(ctx)
  1826  
  1827  		err2 := f2.Get(ctx, nil)
  1828  		if _, ok := err2.(*CanceledError); !ok {
  1829  			return "", err2
  1830  		}
  1831  
  1832  		err := f.Get(ctx, &result)
  1833  		return result, err
  1834  	}
  1835  
  1836  	env := s.NewTestWorkflowEnvironment()
  1837  	env.RegisterWorkflow(workflowFn)
  1838  	env.OnActivity(localActivityFn, mock.Anything, "local_activity").Return("hello mock", nil).Once()
  1839  	var startedCount, completedCount, canceledCount int
  1840  	env.SetOnLocalActivityStartedListener(func(activityInfo *ActivityInfo, ctx context.Context, args []interface{}) {
  1841  		startedCount++
  1842  	})
  1843  
  1844  	env.SetOnLocalActivityCompletedListener(func(activityInfo *ActivityInfo, result Value, err error) {
  1845  		s.NoError(err)
  1846  		var resultValue string
  1847  		err = result.Get(&resultValue)
  1848  		s.NoError(err)
  1849  		s.Equal("hello mock", resultValue)
  1850  		completedCount++
  1851  	})
  1852  
  1853  	env.SetOnLocalActivityCanceledListener(func(activityInfo *ActivityInfo) {
  1854  		canceledCount++
  1855  	})
  1856  
  1857  	env.ExecuteWorkflow(workflowFn)
  1858  	env.AssertExpectations(s.T())
  1859  	s.Equal(2, startedCount)
  1860  	s.Equal(1, completedCount)
  1861  	s.Equal(1, canceledCount)
  1862  	s.True(env.IsWorkflowCompleted())
  1863  	s.NoError(env.GetWorkflowError())
  1864  	var result string
  1865  	err := env.GetWorkflowResult(&result)
  1866  	s.NoError(err)
  1867  	s.Equal("hello mock", result)
  1868  }
  1869  
  1870  func (s *WorkflowTestSuiteUnitTest) Test_SignalChildWorkflow() {
  1871  	// This test will send signal from parent to child, and then child will send back signal to ack. No mock is needed.
  1872  	signalName := "test-signal-name"
  1873  	signalData := "test-signal-data"
  1874  	childWorkflowFn := func(ctx Context, parentExec WorkflowExecution) (string, error) {
  1875  		var data string
  1876  		GetSignalChannel(ctx, signalName).Receive(ctx, &data)
  1877  
  1878  		err := SignalExternalWorkflow(ctx, parentExec.ID, parentExec.RunID, signalName, data+"-received").Get(ctx, nil)
  1879  		if err != nil {
  1880  			return "", err
  1881  		}
  1882  
  1883  		return data + "-processed", nil
  1884  	}
  1885  
  1886  	workflowFn := func(ctx Context) error {
  1887  		cwo := ChildWorkflowOptions{
  1888  			ExecutionStartToCloseTimeout: time.Minute,
  1889  			Domain:                       "test-domain",
  1890  		}
  1891  		ctx = WithChildWorkflowOptions(ctx, cwo)
  1892  		childFuture := ExecuteChildWorkflow(ctx, childWorkflowFn, GetWorkflowInfo(ctx).WorkflowExecution)
  1893  
  1894  		// send signal to child workflow
  1895  		signalFuture := childFuture.SignalChildWorkflow(ctx, signalName, signalData)
  1896  		if err := signalFuture.Get(ctx, nil); err != nil {
  1897  			return err
  1898  		}
  1899  
  1900  		// receiving ack signal from child
  1901  		c := GetSignalChannel(ctx, signalName)
  1902  		var ackMsg string
  1903  		c.Receive(ctx, &ackMsg)
  1904  		s.Equal(signalData+"-received", ackMsg)
  1905  
  1906  		var childResult string
  1907  		if err := childFuture.Get(ctx, &childResult); err != nil {
  1908  			return err
  1909  		}
  1910  
  1911  		s.Equal(signalData+"-processed", childResult)
  1912  		return nil
  1913  	}
  1914  
  1915  	env := s.NewTestWorkflowEnvironment()
  1916  	env.RegisterWorkflow(childWorkflowFn)
  1917  	env.RegisterWorkflow(workflowFn)
  1918  	env.ExecuteWorkflow(workflowFn)
  1919  	s.True(env.IsWorkflowCompleted())
  1920  	s.NoError(env.GetWorkflowError())
  1921  }
  1922  
  1923  func (s *WorkflowTestSuiteUnitTest) Test_SignalExternalWorkflow() {
  1924  	signalName := "test-signal-name"
  1925  	signalData := "test-signal-data"
  1926  	workflowFn := func(ctx Context) error {
  1927  		// set domain to be more specific
  1928  		ctx = WithWorkflowDomain(ctx, "test-domain")
  1929  		f1 := SignalExternalWorkflow(ctx, "test-workflow-id1", "test-runid1", signalName, signalData)
  1930  		f2 := SignalExternalWorkflow(ctx, "test-workflow-id2", "test-runid2", signalName, signalData)
  1931  		f3 := SignalExternalWorkflow(ctx, "test-workflow-id3", "test-runid3", signalName, signalData)
  1932  
  1933  		// signal1 succeed
  1934  		err1 := f1.Get(ctx, nil)
  1935  		if err1 != nil {
  1936  			return err1
  1937  		}
  1938  
  1939  		// signal2 failed
  1940  		err2 := f2.Get(ctx, nil)
  1941  		if err2 == nil {
  1942  			return errors.New("signal2 should fail")
  1943  		}
  1944  
  1945  		// signal3 succeed with delay
  1946  		t := NewTimer(ctx, time.Second*30)
  1947  		timerFired, signalSend := false, false
  1948  
  1949  		NewSelector(ctx).AddFuture(t, func(f Future) {
  1950  			timerFired = true
  1951  		}).AddFuture(f3, func(f Future) {
  1952  			signalSend = true
  1953  		}).Select(ctx)
  1954  
  1955  		// verify that timer fired first because the signal will delay 1 minute
  1956  		s.True(timerFired)
  1957  		s.False(signalSend)
  1958  
  1959  		err3 := f3.Get(ctx, nil)
  1960  		if err3 != nil {
  1961  			return err3
  1962  		}
  1963  
  1964  		return nil
  1965  	}
  1966  
  1967  	env := s.NewTestWorkflowEnvironment()
  1968  	env.RegisterWorkflow(workflowFn)
  1969  
  1970  	// signal1 should succeed
  1971  	env.OnSignalExternalWorkflow("test-domain", "test-workflow-id1", "test-runid1", signalName, signalData).Return(nil).Once()
  1972  
  1973  	// signal2 should fail
  1974  	env.OnSignalExternalWorkflow("test-domain", "test-workflow-id2", "test-runid2", signalName, signalData).Return(
  1975  		func(domainName, workflowID, runID, signalName string, arg interface{}) error {
  1976  			return errors.New("unknown external workflow")
  1977  		}).Once()
  1978  
  1979  	// signal3 should succeed with delay, mock match exactly the parameters
  1980  	env.OnSignalExternalWorkflow("test-domain", "test-workflow-id3", "test-runid3", signalName, signalData).After(time.Minute).Return(nil).Once()
  1981  
  1982  	env.ExecuteWorkflow(workflowFn)
  1983  	env.AssertExpectations(s.T())
  1984  	s.True(env.IsWorkflowCompleted())
  1985  	s.NoError(env.GetWorkflowError())
  1986  }
  1987  
  1988  func (s *WorkflowTestSuiteUnitTest) Test_CancelChildWorkflow() {
  1989  	childWorkflowFn := func(ctx Context) error {
  1990  		var err error
  1991  		selector := NewSelector(ctx)
  1992  		timer := NewTimer(ctx, 10*time.Second)
  1993  		selector.AddFuture(timer, func(f Future) {
  1994  			err = f.Get(ctx, nil)
  1995  		}).Select(ctx)
  1996  
  1997  		GetLogger(ctx).Info("child workflow returned error", zap.Error(err))
  1998  		return err
  1999  	}
  2000  
  2001  	workflowFn := func(ctx Context) error {
  2002  
  2003  		cwo := ChildWorkflowOptions{
  2004  			Domain:                       "test-domain",
  2005  			ExecutionStartToCloseTimeout: time.Minute,
  2006  		}
  2007  
  2008  		childCtx := WithChildWorkflowOptions(ctx, cwo)
  2009  		childCtx, cancel := WithCancel(childCtx)
  2010  		childFuture := ExecuteChildWorkflow(childCtx, childWorkflowFn)
  2011  		Sleep(ctx, 2*time.Second)
  2012  		cancel()
  2013  
  2014  		err := childFuture.Get(childCtx, nil)
  2015  		if _, ok := err.(*CanceledError); !ok {
  2016  			return fmt.Errorf("Cancel child workflow should receive CanceledError, instead got: %v", err)
  2017  		}
  2018  		return nil
  2019  	}
  2020  
  2021  	env := s.NewTestWorkflowEnvironment()
  2022  	env.RegisterWorkflow(childWorkflowFn)
  2023  	env.RegisterWorkflow(workflowFn)
  2024  	env.ExecuteWorkflow(workflowFn)
  2025  	s.True(env.IsWorkflowCompleted())
  2026  	s.NoError(env.GetWorkflowError())
  2027  }
  2028  
  2029  func (s *WorkflowTestSuiteUnitTest) Test_CancelExternalWorkflow() {
  2030  	workflowFn := func(ctx Context) error {
  2031  		// set domain to be more specific
  2032  		ctx = WithWorkflowDomain(ctx, "test-domain")
  2033  		f1 := RequestCancelExternalWorkflow(ctx, "test-workflow-id1", "test-runid1")
  2034  		f2 := RequestCancelExternalWorkflow(ctx, "test-workflow-id2", "test-runid2")
  2035  
  2036  		// cancellation 1 succeed
  2037  		err1 := f1.Get(ctx, nil)
  2038  		if err1 != nil {
  2039  			return err1
  2040  		}
  2041  
  2042  		// cancellation 2 failed
  2043  		err2 := f2.Get(ctx, nil)
  2044  		if err2 == nil {
  2045  			return errors.New("cancellation 2 should fail")
  2046  		}
  2047  
  2048  		return nil
  2049  	}
  2050  
  2051  	env := s.NewTestWorkflowEnvironment()
  2052  	env.RegisterWorkflow(workflowFn)
  2053  
  2054  	// cancellation 1 should succeed
  2055  	env.OnRequestCancelExternalWorkflow("test-domain", "test-workflow-id1", "test-runid1").Return(nil).Once()
  2056  
  2057  	// cancellation 2 should fail
  2058  	env.OnRequestCancelExternalWorkflow("test-domain", "test-workflow-id2", "test-runid2").Return(
  2059  		func(domainName, workflowID, runID string) error {
  2060  			return errors.New("unknown external workflow")
  2061  		}).Once()
  2062  
  2063  	env.ExecuteWorkflow(workflowFn)
  2064  	env.AssertExpectations(s.T())
  2065  	s.True(env.IsWorkflowCompleted())
  2066  	s.NoError(env.GetWorkflowError())
  2067  }
  2068  
  2069  func (s *WorkflowTestSuiteUnitTest) Test_DisconnectedContext() {
  2070  	childWorkflowFn := func(ctx Context) (string, error) {
  2071  		err := NewTimer(ctx, time.Minute*10).Get(ctx, nil)
  2072  		if _, ok := err.(*CanceledError); ok {
  2073  			dCtx, _ := NewDisconnectedContext(ctx)
  2074  			dCtx = WithActivityOptions(dCtx, s.activityOptions)
  2075  			var cleanupResult string
  2076  			err := ExecuteActivity(dCtx, testActivityHello, "cleanup").Get(dCtx, &cleanupResult)
  2077  			return cleanupResult, err
  2078  		}
  2079  
  2080  		// unexpected
  2081  		return "", errors.New("should not reach here")
  2082  	}
  2083  
  2084  	workflowFn := func(ctx Context) (string, error) {
  2085  		cwo := ChildWorkflowOptions{
  2086  			ExecutionStartToCloseTimeout: time.Hour,
  2087  			WaitForCancellation:          true,
  2088  		}
  2089  		ctx = WithChildWorkflowOptions(ctx, cwo)
  2090  		childCtx, cancelChild := WithCancel(ctx)
  2091  		childFuture := ExecuteChildWorkflow(childCtx, childWorkflowFn) // execute child workflow which in turn execute another child
  2092  		NewSelector(ctx).AddFuture(childFuture, func(f Future) {
  2093  			s.Fail("f1 should not complete before t1")
  2094  		}).AddFuture(NewTimer(ctx, time.Minute), func(f Future) {
  2095  			cancelChild() // child workflow takes too long, cancel child workflow
  2096  		}).Select(ctx)
  2097  
  2098  		var result string
  2099  		err := childFuture.Get(childCtx, &result)
  2100  		return result, err
  2101  	}
  2102  
  2103  	env := s.NewTestWorkflowEnvironment()
  2104  	env.RegisterWorkflow(childWorkflowFn)
  2105  	env.RegisterWorkflow(workflowFn)
  2106  	env.RegisterActivity(testActivityHello)
  2107  
  2108  	env.ExecuteWorkflow(workflowFn)
  2109  
  2110  	s.True(env.IsWorkflowCompleted())
  2111  	s.NoError(env.GetWorkflowError())
  2112  	var workflowResult string
  2113  	err := env.GetWorkflowResult(&workflowResult)
  2114  	s.NoError(err)
  2115  	s.Equal("hello_cleanup", workflowResult)
  2116  }
  2117  
  2118  func (s *WorkflowTestSuiteUnitTest) Test_WorkflowIDReusePolicy() {
  2119  	workflowFn := func(ctx Context) (string, error) {
  2120  		cwo := ChildWorkflowOptions{
  2121  			ExecutionStartToCloseTimeout: time.Minute,
  2122  			WorkflowID:                   "test-child-workflow-id",
  2123  			WorkflowIDReusePolicy:        WorkflowIDReusePolicyRejectDuplicate,
  2124  		}
  2125  		ctx = WithChildWorkflowOptions(ctx, cwo)
  2126  		var helloWorkflowResult string
  2127  		f := ExecuteChildWorkflow(ctx, testWorkflowHello)
  2128  		err := f.GetChildWorkflowExecution().Get(ctx, nil)
  2129  		s.NoError(err)
  2130  		err = f.Get(ctx, &helloWorkflowResult)
  2131  		s.NoError(err)
  2132  
  2133  		// start child with duplicate workflow id, but with policy that won't allow duplicate
  2134  		f = ExecuteChildWorkflow(ctx, testWorkflowHello)
  2135  		err = f.GetChildWorkflowExecution().Get(ctx, nil)
  2136  		s.Error(err)
  2137  		err = f.Get(ctx, &helloWorkflowResult)
  2138  		s.Error(err)
  2139  
  2140  		// now with policy allow duplicate
  2141  		cwo.WorkflowIDReusePolicy = WorkflowIDReusePolicyAllowDuplicate
  2142  		ctx = WithChildWorkflowOptions(ctx, cwo)
  2143  		f = ExecuteChildWorkflow(ctx, testWorkflowHello)
  2144  		err = f.GetChildWorkflowExecution().Get(ctx, nil)
  2145  		s.NoError(err)
  2146  		err = f.Get(ctx, &helloWorkflowResult)
  2147  		s.NoError(err)
  2148  
  2149  		return helloWorkflowResult, nil
  2150  	}
  2151  
  2152  	env := s.NewTestWorkflowEnvironment()
  2153  	env.RegisterWorkflow(workflowFn)
  2154  	env.RegisterWorkflow(testWorkflowHello)
  2155  	env.RegisterActivity(testActivityHello)
  2156  	env.ExecuteWorkflow(workflowFn)
  2157  	var actualResult string
  2158  	s.NoError(env.GetWorkflowResult(&actualResult))
  2159  	s.Equal("hello_world", actualResult)
  2160  }
  2161  
  2162  func (s *WorkflowTestSuiteUnitTest) Test_Channel() {
  2163  	workflowFn := func(ctx Context) error {
  2164  
  2165  		signalCh := GetSignalChannel(ctx, "test-signal")
  2166  		doneCh := NewBufferedChannel(ctx, 100)
  2167  		selector := NewSelector(ctx)
  2168  
  2169  		selector.AddReceive(signalCh, func(c Channel, more bool) {
  2170  		}).AddReceive(doneCh, func(c Channel, more bool) {
  2171  			var doneSignal string
  2172  			c.Receive(ctx, &doneSignal)
  2173  		})
  2174  
  2175  		fanoutChs := []Channel{NewBufferedChannel(ctx, 100), NewBufferedChannel(ctx, 100)}
  2176  
  2177  		processedCount := 0
  2178  		runningCount := 0
  2179  
  2180  	mainLoop:
  2181  		for {
  2182  			selector.Select(ctx)
  2183  			var signal string
  2184  			if !signalCh.ReceiveAsync(&signal) {
  2185  				if runningCount > 0 {
  2186  					continue mainLoop
  2187  				}
  2188  
  2189  				if processedCount < 4 {
  2190  					continue mainLoop
  2191  				}
  2192  
  2193  				// continue as new
  2194  				return NewContinueAsNewError(ctx, "this-workflow")
  2195  			}
  2196  
  2197  			for i := range fanoutChs {
  2198  				ch := fanoutChs[i]
  2199  				ch.SendAsync(signal)
  2200  				processedCount++
  2201  				runningCount++
  2202  				Go(ctx, func(ctx Context) {
  2203  					doneCh.SendAsync("done")
  2204  					runningCount--
  2205  				})
  2206  			}
  2207  		}
  2208  
  2209  		return nil
  2210  	}
  2211  
  2212  	env := s.NewTestWorkflowEnvironment()
  2213  	env.RegisterWorkflow(workflowFn)
  2214  
  2215  	env.RegisterDelayedCallback(func() {
  2216  		env.SignalWorkflow("test-signal", "s1")
  2217  		env.SignalWorkflow("test-signal", "s2")
  2218  	}, time.Minute)
  2219  
  2220  	env.ExecuteWorkflow(workflowFn)
  2221  
  2222  	s.True(env.IsWorkflowCompleted())
  2223  	s.Error(env.GetWorkflowError())
  2224  	_, ok := env.GetWorkflowError().(*ContinueAsNewError)
  2225  	s.True(ok)
  2226  }
  2227  
  2228  func (s *WorkflowTestSuiteUnitTest) Test_ContextMisuse() {
  2229  	workflowFn := func(ctx Context) error {
  2230  		ch := NewChannel(ctx)
  2231  
  2232  		Go(ctx, func(shouldUseThisCtx Context) {
  2233  			Sleep(ctx, time.Hour)
  2234  			ch.Send(ctx, "done")
  2235  		})
  2236  
  2237  		var done string
  2238  		ch.Receive(ctx, &done)
  2239  
  2240  		return nil
  2241  	}
  2242  
  2243  	env := s.NewTestWorkflowEnvironment()
  2244  	env.RegisterWorkflow(workflowFn)
  2245  	env.ExecuteWorkflow(workflowFn)
  2246  
  2247  	s.True(env.IsWorkflowCompleted())
  2248  	s.Error(env.GetWorkflowError())
  2249  	s.Contains(env.GetWorkflowError().Error(), "block on coroutine which is already blocked")
  2250  }
  2251  
  2252  func (s *WorkflowTestSuiteUnitTest) Test_DrainSignalChannel() {
  2253  	workflowFn := func(ctx Context) (string, error) {
  2254  
  2255  		signalCh := GetSignalChannel(ctx, "test-signal")
  2256  		var s1, s2, s3 string
  2257  		signalCh.Receive(ctx, &s1)
  2258  		if !signalCh.ReceiveAsync(&s2) {
  2259  			return "", errors.New("expected signal")
  2260  		}
  2261  		if signalCh.ReceiveAsync(&s3) {
  2262  			return "", errors.New("unexpected signal")
  2263  		}
  2264  		return s1 + s2, nil
  2265  	}
  2266  
  2267  	env := s.NewTestWorkflowEnvironment()
  2268  	env.RegisterWorkflow(workflowFn)
  2269  
  2270  	env.RegisterDelayedCallback(func() {
  2271  		env.SignalWorkflowSkippingDecision("test-signal", "s1")
  2272  		env.SignalWorkflow("test-signal", "s2")
  2273  	}, time.Minute)
  2274  
  2275  	env.ExecuteWorkflow(workflowFn)
  2276  
  2277  	s.True(env.IsWorkflowCompleted())
  2278  	s.NoError(env.GetWorkflowError())
  2279  	var result string
  2280  	env.GetWorkflowResult(&result)
  2281  	s.Equal("s1s2", result)
  2282  }
  2283  
  2284  func (s *WorkflowTestSuiteUnitTest) Test_ActivityRetry() {
  2285  	attempt1Count := 0
  2286  	activityFailedFn := func(ctx context.Context) (string, error) {
  2287  		attempt1Count++
  2288  		return "", NewCustomError("bad-bug")
  2289  	}
  2290  
  2291  	attempt2Count := 0
  2292  	activityFn := func(ctx context.Context) (string, error) {
  2293  		attempt2Count++
  2294  		info := GetActivityInfo(ctx)
  2295  		if info.Attempt < 2 {
  2296  			return "", NewCustomError("bad-luck")
  2297  		}
  2298  		return "retry-done", nil
  2299  	}
  2300  
  2301  	workflowFn := func(ctx Context) (string, error) {
  2302  		ao := ActivityOptions{
  2303  			ScheduleToStartTimeout: time.Minute,
  2304  			StartToCloseTimeout:    time.Minute,
  2305  			RetryPolicy: &RetryPolicy{
  2306  				MaximumAttempts:          5,
  2307  				InitialInterval:          time.Second,
  2308  				MaximumInterval:          time.Second * 10,
  2309  				BackoffCoefficient:       2,
  2310  				NonRetriableErrorReasons: []string{"bad-bug"},
  2311  				ExpirationInterval:       time.Minute,
  2312  			},
  2313  		}
  2314  		ctx = WithActivityOptions(ctx, ao)
  2315  
  2316  		err := ExecuteActivity(ctx, activityFailedFn).Get(ctx, nil)
  2317  		badBug, ok := err.(*CustomError)
  2318  		s.True(ok)
  2319  		s.Equal("bad-bug", badBug.Reason())
  2320  
  2321  		var result string
  2322  		err = ExecuteActivity(ctx, activityFn).Get(ctx, &result)
  2323  		if err != nil {
  2324  			return "", err
  2325  		}
  2326  		return result, nil
  2327  	}
  2328  
  2329  	env := s.NewTestWorkflowEnvironment()
  2330  	env.RegisterWorkflow(workflowFn)
  2331  	env.RegisterActivity(activityFailedFn)
  2332  	env.RegisterActivity(activityFn)
  2333  
  2334  	// set a workflow timeout timer to test
  2335  	// if the timer will fire during activity retry
  2336  	env.SetWorkflowTimeout(10 * time.Second)
  2337  	env.ExecuteWorkflow(workflowFn)
  2338  
  2339  	s.True(env.IsWorkflowCompleted())
  2340  	s.NoError(env.GetWorkflowError())
  2341  	var result string
  2342  	s.NoError(env.GetWorkflowResult(&result))
  2343  	s.Equal("retry-done", result)
  2344  	s.Equal(1, attempt1Count)
  2345  	s.Equal(3, attempt2Count)
  2346  }
  2347  
  2348  func (s *WorkflowTestSuiteUnitTest) Test_ActivityHeartbeatRetry() {
  2349  	var startedFrom []int
  2350  	activityHeartBeatFn := func(ctx context.Context, firstTaskID, taskCount int) error {
  2351  		i := firstTaskID
  2352  		if HasHeartbeatDetails(ctx) {
  2353  			var lastProcessed int
  2354  			if err := GetHeartbeatDetails(ctx, &lastProcessed); err == nil {
  2355  				i = lastProcessed + 1
  2356  			}
  2357  		}
  2358  
  2359  		startedFrom = append(startedFrom, i)
  2360  
  2361  		for j := 0; i < firstTaskID+taskCount; i, j = i+1, j+1 {
  2362  			// process task i
  2363  			RecordActivityHeartbeat(ctx, i)
  2364  			if j == 2 && i < firstTaskID+taskCount-1 { // simulate failure after processing 3 tasks
  2365  				return NewCustomError("bad-luck")
  2366  			}
  2367  		}
  2368  
  2369  		return nil
  2370  	}
  2371  
  2372  	workflowFn := func(ctx Context) error {
  2373  		ao := ActivityOptions{
  2374  			ScheduleToStartTimeout: time.Minute,
  2375  			StartToCloseTimeout:    time.Minute,
  2376  			RetryPolicy: &RetryPolicy{
  2377  				MaximumAttempts:          3,
  2378  				InitialInterval:          time.Second,
  2379  				MaximumInterval:          time.Second * 10,
  2380  				BackoffCoefficient:       2,
  2381  				NonRetriableErrorReasons: []string{"bad-bug"},
  2382  				ExpirationInterval:       time.Minute,
  2383  			},
  2384  		}
  2385  		ctx = WithActivityOptions(ctx, ao)
  2386  
  2387  		err := ExecuteActivity(ctx, activityHeartBeatFn, 0, 9).Get(ctx, nil)
  2388  		if err != nil {
  2389  			return err
  2390  		}
  2391  		return nil
  2392  	}
  2393  
  2394  	env := s.NewTestWorkflowEnvironment()
  2395  	env.SetTestTimeout(time.Hour)
  2396  	env.RegisterWorkflow(workflowFn)
  2397  	env.RegisterActivity(activityHeartBeatFn)
  2398  	env.ExecuteWorkflow(workflowFn)
  2399  
  2400  	s.True(env.IsWorkflowCompleted())
  2401  	s.NoError(env.GetWorkflowError())
  2402  	s.Equal(3, len(startedFrom))
  2403  	s.Equal([]int{0, 3, 6}, startedFrom)
  2404  }
  2405  
  2406  func (s *WorkflowTestSuiteUnitTest) Test_LocalActivityRetry() {
  2407  	nonretriableCount := 0
  2408  	nonretriableFn := func(ctx context.Context) (string, error) {
  2409  		nonretriableCount++
  2410  		return "", NewCustomError("bad-bug")
  2411  	}
  2412  
  2413  	retriableCount := 0
  2414  	retriableFn := func(ctx context.Context) (string, error) {
  2415  		retriableCount++
  2416  		info := GetActivityInfo(ctx)
  2417  		if info.Attempt < 2 {
  2418  			return "", NewCustomError("bad-luck")
  2419  		}
  2420  		return "retry-done", nil
  2421  	}
  2422  
  2423  	workflowFn := func(ctx Context) (string, error) {
  2424  		lao := LocalActivityOptions{
  2425  			ScheduleToCloseTimeout: time.Minute,
  2426  			RetryPolicy: &RetryPolicy{
  2427  				MaximumAttempts:          3,
  2428  				InitialInterval:          time.Second,
  2429  				MaximumInterval:          time.Second * 10,
  2430  				BackoffCoefficient:       2,
  2431  				NonRetriableErrorReasons: []string{"bad-bug"},
  2432  				ExpirationInterval:       time.Minute,
  2433  			},
  2434  		}
  2435  		ctx = WithLocalActivityOptions(ctx, lao)
  2436  
  2437  		err := ExecuteLocalActivity(ctx, nonretriableFn).Get(ctx, nil)
  2438  		badBug, ok := err.(*CustomError)
  2439  		s.True(ok)
  2440  		s.Equal("bad-bug", badBug.Reason())
  2441  
  2442  		var result string
  2443  		err = ExecuteLocalActivity(ctx, retriableFn).Get(ctx, &result)
  2444  		if err != nil {
  2445  			return "", err
  2446  		}
  2447  		return result, nil
  2448  	}
  2449  
  2450  	env := s.NewTestWorkflowEnvironment()
  2451  	env.RegisterWorkflow(workflowFn)
  2452  	env.ExecuteWorkflow(workflowFn)
  2453  
  2454  	s.True(env.IsWorkflowCompleted())
  2455  	s.NoError(env.GetWorkflowError())
  2456  	var result string
  2457  	s.NoError(env.GetWorkflowResult(&result))
  2458  	s.Equal("retry-done", result)
  2459  	s.Equal(1, nonretriableCount)
  2460  	s.Equal(3, retriableCount)
  2461  }
  2462  
  2463  func (s *WorkflowTestSuiteUnitTest) Test_LocalActivityRetryOnCancel() {
  2464  	attempts := 0
  2465  	localActivityFn := func(ctx context.Context) (int32, error) {
  2466  		attempts++
  2467  		info := GetActivityInfo(ctx)
  2468  		if info.Attempt < 2 {
  2469  			return int32(-1), NewCanceledError("details")
  2470  		}
  2471  		return info.Attempt, nil
  2472  	}
  2473  
  2474  	workflowFn := func(ctx Context) (int32, error) {
  2475  		lao := LocalActivityOptions{
  2476  			ScheduleToCloseTimeout: time.Minute,
  2477  			RetryPolicy: &RetryPolicy{
  2478  				MaximumAttempts:          3,
  2479  				InitialInterval:          time.Second,
  2480  				MaximumInterval:          time.Second * 10,
  2481  				BackoffCoefficient:       2,
  2482  				NonRetriableErrorReasons: []string{"bad-bug"},
  2483  				ExpirationInterval:       time.Minute,
  2484  			},
  2485  		}
  2486  		ctx = WithLocalActivityOptions(ctx, lao)
  2487  
  2488  		var result int32
  2489  		err := ExecuteLocalActivity(ctx, localActivityFn).Get(ctx, &result)
  2490  		if err != nil {
  2491  			return int32(-1), err
  2492  		}
  2493  		return result, nil
  2494  	}
  2495  
  2496  	env := s.NewTestWorkflowEnvironment()
  2497  	env.RegisterWorkflow(workflowFn)
  2498  	env.ExecuteWorkflow(workflowFn)
  2499  
  2500  	s.True(env.IsWorkflowCompleted())
  2501  	s.Error(env.GetWorkflowError())
  2502  	s.True(IsCanceledError(env.GetWorkflowError()))
  2503  	s.Equal(1, attempts)
  2504  }
  2505  
  2506  func (s *WorkflowTestSuiteUnitTest) Test_ActivityRetryOnCancel() {
  2507  	workflowFn := func(ctx Context) (int32, error) {
  2508  		ao := ActivityOptions{
  2509  			ScheduleToStartTimeout: time.Minute,
  2510  			StartToCloseTimeout:    time.Minute,
  2511  			RetryPolicy: &RetryPolicy{
  2512  				MaximumAttempts:          3,
  2513  				InitialInterval:          time.Second,
  2514  				MaximumInterval:          time.Second * 10,
  2515  				BackoffCoefficient:       2,
  2516  				NonRetriableErrorReasons: []string{"bad-bug"},
  2517  				ExpirationInterval:       time.Minute,
  2518  			},
  2519  		}
  2520  		ctx = WithActivityOptions(ctx, ao)
  2521  
  2522  		var result int32
  2523  		err := ExecuteActivity(ctx, testActivityCanceled).Get(ctx, &result)
  2524  		if err != nil {
  2525  			return int32(-1), err
  2526  		}
  2527  		return result, nil
  2528  	}
  2529  
  2530  	env := s.NewTestWorkflowEnvironment()
  2531  	env.RegisterWorkflow(workflowFn)
  2532  	env.RegisterActivity(testActivityCanceled)
  2533  
  2534  	env.ExecuteWorkflow(workflowFn)
  2535  
  2536  	s.True(env.IsWorkflowCompleted())
  2537  	s.Error(env.GetWorkflowError())
  2538  	s.True(IsCanceledError(env.GetWorkflowError()))
  2539  }
  2540  
  2541  func (s *WorkflowTestSuiteUnitTest) Test_ChildWorkflowRetry() {
  2542  
  2543  	childWorkflowFn := func(ctx Context) (string, error) {
  2544  		info := GetWorkflowInfo(ctx)
  2545  		if info.Attempt < 2 {
  2546  			return "", NewCustomError("bad-luck")
  2547  		}
  2548  		return "retry-done", nil
  2549  	}
  2550  
  2551  	workflowFn := func(ctx Context) (string, error) {
  2552  		cwo := ChildWorkflowOptions{
  2553  			ExecutionStartToCloseTimeout: time.Minute,
  2554  			RetryPolicy: &RetryPolicy{
  2555  				MaximumAttempts:          3,
  2556  				InitialInterval:          time.Second,
  2557  				MaximumInterval:          time.Second * 10,
  2558  				BackoffCoefficient:       2,
  2559  				NonRetriableErrorReasons: []string{"bad-bug"},
  2560  				ExpirationInterval:       time.Minute,
  2561  			},
  2562  		}
  2563  		ctx = WithChildWorkflowOptions(ctx, cwo)
  2564  		var childResult string
  2565  		err := ExecuteChildWorkflow(ctx, childWorkflowFn).Get(ctx, &childResult)
  2566  		if err != nil {
  2567  			return "", err
  2568  		}
  2569  
  2570  		return childResult, nil
  2571  	}
  2572  
  2573  	env := s.NewTestWorkflowEnvironment()
  2574  	env.RegisterWorkflow(childWorkflowFn)
  2575  	env.RegisterWorkflow(workflowFn)
  2576  	env.ExecuteWorkflow(workflowFn)
  2577  
  2578  	s.True(env.IsWorkflowCompleted())
  2579  	s.NoError(env.GetWorkflowError())
  2580  	var result string
  2581  	s.NoError(env.GetWorkflowResult(&result))
  2582  	s.Equal("retry-done", result)
  2583  }
  2584  
  2585  func (s *WorkflowTestSuiteUnitTest) Test_SignalChildWorkflowRetry() {
  2586  	childWorkflowFn := func(ctx Context) (string, error) {
  2587  		info := GetWorkflowInfo(ctx)
  2588  		if info.Attempt < 2 {
  2589  			return "", NewCustomError("bad-luck")
  2590  		}
  2591  
  2592  		ch := GetSignalChannel(ctx, "test-signal-name")
  2593  		timeout := NewTimer(ctx, time.Second*3)
  2594  		s := NewSelector(ctx)
  2595  		var signal string
  2596  		s.AddFuture(timeout, func(f Future) {
  2597  			signal = "timeout"
  2598  		}).AddReceive(ch, func(c Channel, more bool) {
  2599  			c.Receive(ctx, &signal)
  2600  		}).Select(ctx)
  2601  
  2602  		return signal, nil
  2603  	}
  2604  
  2605  	workflowFn := func(ctx Context) (string, error) {
  2606  		cwo := ChildWorkflowOptions{
  2607  			WorkflowID:                   "test-retry-signal-child-workflow",
  2608  			ExecutionStartToCloseTimeout: time.Minute,
  2609  			RetryPolicy: &RetryPolicy{
  2610  				MaximumAttempts:          3,
  2611  				InitialInterval:          time.Second * 3,
  2612  				MaximumInterval:          time.Second * 3,
  2613  				BackoffCoefficient:       1,
  2614  				NonRetriableErrorReasons: []string{"bad-bug"},
  2615  				ExpirationInterval:       time.Minute,
  2616  			},
  2617  		}
  2618  		ctx = WithChildWorkflowOptions(ctx, cwo)
  2619  		var childResult string
  2620  		err := ExecuteChildWorkflow(ctx, childWorkflowFn).Get(ctx, &childResult)
  2621  		if err != nil {
  2622  			return "", err
  2623  		}
  2624  
  2625  		return childResult, nil
  2626  	}
  2627  
  2628  	env := s.NewTestWorkflowEnvironment()
  2629  	env.RegisterWorkflow(childWorkflowFn)
  2630  	env.RegisterWorkflow(workflowFn)
  2631  
  2632  	env.RegisterDelayedCallback(func() {
  2633  		env.SignalWorkflowByID("test-retry-signal-child-workflow", "test-signal-name", "test-signal-data")
  2634  	}, time.Second*7 /* after 2nd attempt failed, but before 3rd attempt starts */)
  2635  
  2636  	env.ExecuteWorkflow(workflowFn)
  2637  
  2638  	s.True(env.IsWorkflowCompleted())
  2639  	s.NoError(env.GetWorkflowError())
  2640  	var result string
  2641  	s.NoError(env.GetWorkflowResult(&result))
  2642  	s.Equal("test-signal-data", result)
  2643  }
  2644  
  2645  func (s *WorkflowTestSuiteUnitTest) Test_TestWorkflowTimeoutInBusyLoop() {
  2646  	neverEndingWorkflow := func(ctx Context) error {
  2647  		for {
  2648  			Sleep(ctx, time.Hour)
  2649  		}
  2650  	}
  2651  
  2652  	env := s.NewTestWorkflowEnvironment()
  2653  	env.SetWorkflowTimeout((time.Hour * 10) + time.Minute)
  2654  	timerFiredCount := 0
  2655  	env.SetOnTimerFiredListener(func(timerID string) {
  2656  		timerFiredCount++
  2657  	})
  2658  	env.RegisterWorkflow(neverEndingWorkflow)
  2659  
  2660  	env.ExecuteWorkflow(neverEndingWorkflow)
  2661  	s.Equal(10, timerFiredCount)
  2662  	s.Error(env.GetWorkflowError())
  2663  	_, ok := env.GetWorkflowError().(*TimeoutError)
  2664  	s.True(ok)
  2665  }
  2666  
  2667  func (s *WorkflowTestSuiteUnitTest) Test_TestChildWorkflowTimeout() {
  2668  	childWorkflowFn := func(ctx Context) error {
  2669  		Sleep(ctx, time.Hour*5)
  2670  		return nil
  2671  	}
  2672  
  2673  	workflowFn := func(ctx Context) error {
  2674  		cwo := ChildWorkflowOptions{
  2675  			ExecutionStartToCloseTimeout: time.Hour * 3, // less than 5h that child workflow would take.
  2676  		}
  2677  		ctx = WithChildWorkflowOptions(ctx, cwo)
  2678  		err := ExecuteChildWorkflow(ctx, childWorkflowFn).Get(ctx, nil)
  2679  
  2680  		s.Error(err)
  2681  		if _, ok := err.(*TimeoutError); ok {
  2682  			return nil
  2683  		}
  2684  		return err
  2685  	}
  2686  
  2687  	env := s.NewTestWorkflowEnvironment()
  2688  	env.SetWorkflowTimeout(time.Hour * 10)
  2689  	timerFiredCount := 0
  2690  	env.SetOnTimerFiredListener(func(timerID string) {
  2691  		timerFiredCount++
  2692  	})
  2693  	env.RegisterWorkflow(childWorkflowFn)
  2694  	env.RegisterWorkflow(workflowFn)
  2695  
  2696  	env.ExecuteWorkflow(workflowFn)
  2697  	s.NoError(env.GetWorkflowError())
  2698  }
  2699  
  2700  func (s *WorkflowTestSuiteUnitTest) Test_SameActivityIDFromDifferentChildWorkflow() {
  2701  	childWorkflowFn := func(ctx Context) (string, error) {
  2702  		ao := ActivityOptions{
  2703  			ActivityID:             "per_workflow_unique_activity_id",
  2704  			ScheduleToStartTimeout: time.Minute,
  2705  			StartToCloseTimeout:    time.Minute,
  2706  		}
  2707  		ctx = WithActivityOptions(ctx, ao)
  2708  
  2709  		info := GetWorkflowInfo(ctx)
  2710  
  2711  		var result string
  2712  		err := ExecuteActivity(ctx, testActivityHello, info.WorkflowExecution.ID).Get(ctx, &result)
  2713  		if err != nil {
  2714  			return "", err
  2715  		}
  2716  
  2717  		return result, nil
  2718  	}
  2719  
  2720  	workflowFn := func(ctx Context) (string, error) {
  2721  		ctx1 := WithChildWorkflowOptions(ctx, ChildWorkflowOptions{
  2722  			WorkflowID:                   "child_1",
  2723  			ExecutionStartToCloseTimeout: time.Minute,
  2724  		})
  2725  		f1 := ExecuteChildWorkflow(ctx1, childWorkflowFn)
  2726  
  2727  		ctx2 := WithChildWorkflowOptions(ctx, ChildWorkflowOptions{
  2728  			WorkflowID:                   "child_2",
  2729  			ExecutionStartToCloseTimeout: time.Minute,
  2730  		})
  2731  		f2 := ExecuteChildWorkflow(ctx2, childWorkflowFn)
  2732  
  2733  		var result1, result2 string
  2734  		if err := f1.Get(ctx1, &result1); err != nil {
  2735  			return "", err
  2736  		}
  2737  		if err := f2.Get(ctx1, &result2); err != nil {
  2738  			return "", err
  2739  		}
  2740  
  2741  		return result1 + " " + result2, nil
  2742  	}
  2743  
  2744  	env := s.NewTestWorkflowEnvironment()
  2745  	env.RegisterWorkflow(workflowFn)
  2746  	env.RegisterWorkflow(childWorkflowFn)
  2747  	env.RegisterActivity(testActivityHello)
  2748  
  2749  	env.ExecuteWorkflow(workflowFn)
  2750  
  2751  	s.True(env.IsWorkflowCompleted())
  2752  	s.NoError(env.GetWorkflowError())
  2753  	var actualResult string
  2754  	s.NoError(env.GetWorkflowResult(&actualResult))
  2755  	s.Equal("hello_child_1 hello_child_2", actualResult)
  2756  }
  2757  
  2758  func (s *WorkflowTestSuiteUnitTest) Test_MockChildWorkflowAlreadyRunning() {
  2759  	childWorkflowFn := func(ctx Context) error {
  2760  		return nil
  2761  	}
  2762  
  2763  	runID := "run-id"
  2764  	workflowFn := func(ctx Context) error {
  2765  		cwo := ChildWorkflowOptions{
  2766  			ExecutionStartToCloseTimeout: time.Minute,
  2767  		}
  2768  		ctx = WithChildWorkflowOptions(ctx, cwo)
  2769  		err := ExecuteChildWorkflow(ctx, childWorkflowFn).Get(ctx, nil)
  2770  		s.Error(err)
  2771  
  2772  		alreadySytartedErr, ok := err.(*shared.WorkflowExecutionAlreadyStartedError)
  2773  		s.True(ok)
  2774  		s.Equal(runID, *alreadySytartedErr.RunId)
  2775  
  2776  		return nil
  2777  	}
  2778  
  2779  	env := s.NewTestWorkflowEnvironment()
  2780  	RegisterWorkflow(childWorkflowFn)
  2781  	RegisterWorkflow(workflowFn)
  2782  
  2783  	env.OnWorkflow(childWorkflowFn, mock.Anything).
  2784  		Return(&shared.WorkflowExecutionAlreadyStartedError{
  2785  			RunId: &runID,
  2786  		})
  2787  
  2788  	env.ExecuteWorkflow(workflowFn)
  2789  	s.NoError(env.GetWorkflowError())
  2790  }
  2791  
  2792  func (s *WorkflowTestSuiteUnitTest) Test_ChildWorkflowAlreadyRunning() {
  2793  	workflowFn := func(ctx Context) (string, error) {
  2794  		ctx1 := WithChildWorkflowOptions(ctx, ChildWorkflowOptions{
  2795  			WorkflowID:                   "Test_ChildWorkflowAlreadyRunning",
  2796  			ExecutionStartToCloseTimeout: time.Minute,
  2797  			// WorkflowIDReusePolicy:        WorkflowIDReusePolicyAllowDuplicate,
  2798  		})
  2799  
  2800  		var result1, result2 string
  2801  		err := ExecuteChildWorkflow(ctx1, testWorkflowHeartbeat, "child1", time.Millisecond).Get(ctx1, &result1)
  2802  		s.NoError(err)
  2803  
  2804  		f2 := ExecuteChildWorkflow(ctx1, testWorkflowHeartbeat, "child2", time.Second)
  2805  
  2806  		err = f2.Get(ctx1, &result2)
  2807  		s.Error(err)
  2808  		_, ok := err.(*shared.WorkflowExecutionAlreadyStartedError)
  2809  		s.True(ok)
  2810  
  2811  		return result1 + " " + result2, nil
  2812  	}
  2813  
  2814  	env := s.NewTestWorkflowEnvironment()
  2815  	env.RegisterWorkflow(workflowFn)
  2816  	env.RegisterWorkflow(testWorkflowHeartbeat)
  2817  	env.RegisterActivity(testActivityHeartbeat)
  2818  	env.ExecuteWorkflow(workflowFn)
  2819  
  2820  	s.True(env.IsWorkflowCompleted())
  2821  	err := env.GetWorkflowError()
  2822  	s.NoError(err)
  2823  
  2824  	var result string
  2825  	s.NoError(env.GetWorkflowResult(&result))
  2826  	s.Equal("heartbeat_child1 ", result)
  2827  }
  2828  
  2829  func (s *WorkflowTestSuiteUnitTest) Test_CronChildWorkflow() {
  2830  	failedCount, successCount, lastCompletionResult := 0, 0, 0
  2831  	cronWorkflow := func(ctx Context) (int, error) {
  2832  		info := GetWorkflowInfo(ctx)
  2833  		var result int
  2834  		if HasLastCompletionResult(ctx) {
  2835  			GetLastCompletionResult(ctx, &result)
  2836  		}
  2837  		Sleep(ctx, time.Second*3)
  2838  		if info.Attempt == 0 {
  2839  			failedCount++
  2840  			return 0, errors.New("please-retry")
  2841  		}
  2842  		successCount++
  2843  		result++
  2844  		lastCompletionResult = result
  2845  		return result, nil
  2846  	}
  2847  
  2848  	testWorkflow := func(ctx Context) error {
  2849  		ctx1 := WithChildWorkflowOptions(ctx, ChildWorkflowOptions{
  2850  			ExecutionStartToCloseTimeout: time.Minute * 10,
  2851  			RetryPolicy: &RetryPolicy{
  2852  				MaximumAttempts:          5,
  2853  				InitialInterval:          time.Second,
  2854  				MaximumInterval:          time.Second * 10,
  2855  				BackoffCoefficient:       2,
  2856  				NonRetriableErrorReasons: []string{"bad-bug"},
  2857  				ExpirationInterval:       time.Hour,
  2858  			},
  2859  			CronSchedule: "0 * * * *", // hourly
  2860  		})
  2861  
  2862  		cronFuture := ExecuteChildWorkflow(ctx1, cronWorkflow) // cron never stop so this future won't return
  2863  
  2864  		timeoutTimer := NewTimer(ctx, time.Hour*3)
  2865  		selector := NewSelector(ctx)
  2866  		var err error
  2867  		selector.AddFuture(cronFuture, func(f Future) {
  2868  			err = errors.New("cron workflow returns, this is not expected")
  2869  		}).AddFuture(timeoutTimer, func(f Future) {
  2870  			// err will be nil
  2871  		}).Select(ctx)
  2872  
  2873  		return err
  2874  	}
  2875  
  2876  	env := s.NewTestWorkflowEnvironment()
  2877  	env.RegisterWorkflow(cronWorkflow)
  2878  	env.RegisterWorkflow(testWorkflow)
  2879  
  2880  	startTime, _ := time.Parse(time.RFC3339, "2018-12-20T16:30:00+08:00")
  2881  	env.SetStartTime(startTime)
  2882  	env.ExecuteWorkflow(testWorkflow)
  2883  
  2884  	s.True(env.IsWorkflowCompleted())
  2885  	err := env.GetWorkflowError()
  2886  	s.NoError(err)
  2887  
  2888  	s.Equal(4, failedCount)
  2889  	s.Equal(4, successCount)
  2890  	s.Equal(4, lastCompletionResult)
  2891  }
  2892  
  2893  func (s *WorkflowTestSuiteUnitTest) Test_CronWorkflow() {
  2894  	var totalRuns int
  2895  	cronWorkflow := func(ctx Context) (int, error) {
  2896  		var result int
  2897  		if HasLastCompletionResult(ctx) {
  2898  			GetLastCompletionResult(ctx, &result)
  2899  		}
  2900  		Sleep(ctx, time.Second*3)
  2901  		result++
  2902  		return result, nil
  2903  	}
  2904  
  2905  	env := s.NewTestWorkflowEnvironment()
  2906  	env.SetWorkflowCronSchedule("0 * * * *") // hourly)
  2907  	env.SetWorkflowCronMaxIterations(1)
  2908  	env.RegisterWorkflow(cronWorkflow)
  2909  
  2910  	startTime, _ := time.Parse(time.RFC3339, "2018-12-20T16:30:00+08:00")
  2911  	env.SetStartTime(startTime)
  2912  	env.ExecuteWorkflow(cronWorkflow)
  2913  
  2914  	env.GetWorkflowResult(&totalRuns)
  2915  	s.True(env.IsWorkflowCompleted())
  2916  	err := env.GetWorkflowError()
  2917  	s.NoError(err)
  2918  
  2919  	s.Equal(2, totalRuns)
  2920  }
  2921  
  2922  func (s *WorkflowTestSuiteUnitTest) Test_CronHasLastResult() {
  2923  	cronWorkflow := func(ctx Context) (int, error) {
  2924  		var result int
  2925  		if HasLastCompletionResult(ctx) {
  2926  			GetLastCompletionResult(ctx, &result)
  2927  		}
  2928  
  2929  		return result + 1, nil
  2930  	}
  2931  
  2932  	env := s.NewTestWorkflowEnvironment()
  2933  	env.RegisterWorkflow(cronWorkflow)
  2934  	lastResult := 3
  2935  	env.SetLastCompletionResult(lastResult)
  2936  	env.ExecuteWorkflow(cronWorkflow)
  2937  
  2938  	s.True(env.IsWorkflowCompleted())
  2939  	s.NoError(env.GetWorkflowError())
  2940  
  2941  	var result int
  2942  	err := env.GetWorkflowResult(&result)
  2943  	s.NoError(err)
  2944  
  2945  	s.Equal(lastResult+1, result)
  2946  }
  2947  
  2948  func (s *WorkflowTestSuiteUnitTest) Test_ActivityWithProgress() {
  2949  	activityFn := func(ctx context.Context) (int, error) {
  2950  		var progress int
  2951  		if HasHeartbeatDetails(ctx) {
  2952  			GetHeartbeatDetails(ctx, &progress)
  2953  		}
  2954  
  2955  		return progress + 1, nil
  2956  	}
  2957  
  2958  	env := s.NewTestActivityEnvironment()
  2959  	env.RegisterActivity(activityFn)
  2960  	lastProgress := 3
  2961  	env.SetHeartbeatDetails(lastProgress)
  2962  	result, err := env.ExecuteActivity(activityFn)
  2963  
  2964  	s.NoError(err)
  2965  
  2966  	var newProgress int
  2967  	err = result.Get(&newProgress)
  2968  	s.NoError(err)
  2969  
  2970  	s.Equal(lastProgress+1, newProgress)
  2971  }
  2972  
  2973  func (s *WorkflowTestSuiteUnitTest) Test_ActivityGoexit() {
  2974  	fn := func(ctx context.Context) error {
  2975  		runtime.Goexit() // usually this is called by t.FailNow(), but can't call FailNow here since that would mark the test as failed.
  2976  		return nil
  2977  	}
  2978  
  2979  	wf := func(ctx Context) error {
  2980  		ao := ActivityOptions{
  2981  			ScheduleToStartTimeout: time.Minute,
  2982  			StartToCloseTimeout:    5 * time.Second,
  2983  		}
  2984  		ctx = WithActivityOptions(ctx, ao)
  2985  		err := ExecuteActivity(ctx, fn).Get(ctx, nil)
  2986  		return err
  2987  	}
  2988  
  2989  	env := s.NewTestWorkflowEnvironment()
  2990  	env.RegisterActivity(fn)
  2991  	env.RegisterWorkflow(wf)
  2992  	env.ExecuteWorkflow(wf)
  2993  	err := env.GetWorkflowError()
  2994  	s.EqualError(err, "activity called runtime.Goexit")
  2995  }
  2996  
  2997  func (s *WorkflowTestSuiteUnitTest) Test_SetWorkerStopChannel() {
  2998  	env := newTestWorkflowEnvironmentImpl(&s.WorkflowTestSuite, nil)
  2999  	c := make(chan struct{})
  3000  	env.setWorkerStopChannel(c)
  3001  	s.NotNil(env.workerStopChannel)
  3002  }
  3003  
  3004  func (s *WorkflowTestSuiteUnitTest) Test_ActivityTimeoutWithDetails() {
  3005  	count := 0
  3006  	timeoutFn := func() error {
  3007  		count++
  3008  		return NewTimeoutError(shared.TimeoutTypeStartToClose, testErrorDetails1)
  3009  	}
  3010  
  3011  	timeoutWf := func(ctx Context) error {
  3012  		ao := ActivityOptions{
  3013  			ScheduleToStartTimeout: time.Minute,
  3014  			StartToCloseTimeout:    5 * time.Second,
  3015  			RetryPolicy: &RetryPolicy{
  3016  				InitialInterval:          time.Second,
  3017  				BackoffCoefficient:       1.1,
  3018  				MaximumAttempts:          3,
  3019  				NonRetriableErrorReasons: []string{"cadenceInternal:Timeout START_TO_CLOSE"},
  3020  			},
  3021  		}
  3022  		ctx = WithActivityOptions(ctx, ao)
  3023  		err := ExecuteActivity(ctx, timeoutFn).Get(ctx, nil)
  3024  		return err
  3025  	}
  3026  
  3027  	wfEnv := s.NewTestWorkflowEnvironment()
  3028  	wfEnv.RegisterWorkflow(timeoutWf)
  3029  	wfEnv.RegisterActivity(timeoutFn)
  3030  
  3031  	wfEnv.ExecuteWorkflow(timeoutWf)
  3032  	err := wfEnv.GetWorkflowError()
  3033  	s.Error(err)
  3034  	timeoutErr, ok := err.(*TimeoutError)
  3035  	s.True(ok)
  3036  	s.Equal(shared.TimeoutTypeStartToClose, timeoutErr.TimeoutType())
  3037  	s.True(timeoutErr.HasDetails())
  3038  	var details string
  3039  	err = timeoutErr.Details(&details)
  3040  	s.NoError(err)
  3041  	s.Equal(testErrorDetails1, details)
  3042  	s.Equal(1, count)
  3043  
  3044  	activityEnv := s.NewTestActivityEnvironment()
  3045  	activityEnv.RegisterActivity(timeoutFn)
  3046  
  3047  	_, err = activityEnv.ExecuteActivity(timeoutFn)
  3048  	s.Error(err)
  3049  	timeoutErr, ok = err.(*TimeoutError)
  3050  	s.True(ok)
  3051  	s.Equal(shared.TimeoutTypeStartToClose, timeoutErr.TimeoutType())
  3052  	s.True(timeoutErr.HasDetails())
  3053  	err = timeoutErr.Details(&details)
  3054  	s.NoError(err)
  3055  	s.Equal(testErrorDetails1, details)
  3056  }
  3057  
  3058  func (s *WorkflowTestSuiteUnitTest) Test_ActivityDeadlineExceeded() {
  3059  	timeoutFn := func(ctx context.Context) error {
  3060  		<-ctx.Done()
  3061  		return nil
  3062  	}
  3063  
  3064  	timeoutWf := func(ctx Context) error {
  3065  		ao := ActivityOptions{
  3066  			ScheduleToStartTimeout: time.Minute,
  3067  			StartToCloseTimeout:    1 * time.Second,
  3068  		}
  3069  		ctx = WithActivityOptions(ctx, ao)
  3070  		err := ExecuteActivity(ctx, timeoutFn).Get(ctx, nil)
  3071  		return err
  3072  	}
  3073  
  3074  	wfEnv := s.NewTestWorkflowEnvironment()
  3075  	wfEnv.RegisterActivity(timeoutFn)
  3076  	wfEnv.RegisterWorkflow(timeoutWf)
  3077  	wfEnv.ExecuteWorkflow(timeoutWf)
  3078  	err := wfEnv.GetWorkflowError()
  3079  	s.Error(err)
  3080  	timeoutErr, ok := err.(*TimeoutError)
  3081  	s.True(ok)
  3082  	s.Equal(shared.TimeoutTypeStartToClose, timeoutErr.TimeoutType())
  3083  	s.True(timeoutErr.HasDetails())
  3084  	var details string
  3085  	err = timeoutErr.Details(&details)
  3086  	s.NoError(err)
  3087  	s.Equal("context deadline exceeded", details)
  3088  }
  3089  
  3090  func (s *WorkflowTestSuiteUnitTest) Test_AwaitWithTimeout() {
  3091  	workflowFn := func(ctx Context) (bool, error) {
  3092  		t := NewTimer(ctx, time.Second)
  3093  		value := false
  3094  		err := Await(ctx, func() bool { return t.IsReady() || value })
  3095  		return value, err
  3096  	}
  3097  
  3098  	env := s.NewTestWorkflowEnvironment()
  3099  	env.RegisterWorkflow(workflowFn)
  3100  	env.ExecuteWorkflow(workflowFn)
  3101  	s.True(env.IsWorkflowCompleted())
  3102  	s.NoError(env.GetWorkflowError())
  3103  	result := true
  3104  	_ = env.GetWorkflowResult(&result)
  3105  	s.False(result)
  3106  }
  3107  
  3108  func (s *WorkflowTestSuiteUnitTest) Test_Regression_ExecuteChildWorkflowWithCanceledContext() {
  3109  	// cancelTime of:
  3110  	// - <0 == do not cancel
  3111  	// - 0  == cancel synchronously
  3112  	// - >0 == cancel after waiting that long
  3113  	check := func(cancelTime time.Duration, bugport bool, expected string) {
  3114  		env := s.NewTestWorkflowEnvironment()
  3115  		env.Test(s.T())
  3116  		env.RegisterWorkflowWithOptions(func(ctx Context) error {
  3117  			return Sleep(ctx, time.Minute)
  3118  		}, RegisterWorkflowOptions{Name: "child"})
  3119  		env.RegisterWorkflowWithOptions(func(ctx Context) (string, error) {
  3120  			ctx, cancel := WithCancel(ctx)
  3121  			if cancelTime == 0 {
  3122  				cancel()
  3123  			} else if cancelTime > 0 {
  3124  				Go(ctx, func(ctx Context) {
  3125  					_ = Sleep(ctx, cancelTime)
  3126  					cancel()
  3127  				})
  3128  			}
  3129  
  3130  			ctx = WithChildWorkflowOptions(ctx, ChildWorkflowOptions{
  3131  				ExecutionStartToCloseTimeout: 2 * time.Minute,
  3132  				TaskStartToCloseTimeout:      2 * time.Minute,
  3133  				Bugports: Bugports{
  3134  					StartChildWorkflowsOnCanceledContext: bugport,
  3135  				},
  3136  			})
  3137  			err := ExecuteChildWorkflow(ctx, "child").Get(ctx, nil)
  3138  
  3139  			if err == nil {
  3140  				return "no err", nil
  3141  			} else if _, ok := err.(*CanceledError); ok {
  3142  				return "canceled", nil
  3143  			}
  3144  			return "unknown: " + err.Error(), nil
  3145  		}, RegisterWorkflowOptions{Name: "parent"})
  3146  
  3147  		env.ExecuteWorkflow("parent")
  3148  		s.True(env.IsWorkflowCompleted())
  3149  		s.NoError(env.GetWorkflowError())
  3150  
  3151  		var result string
  3152  		s.NoError(env.GetWorkflowResult(&result))
  3153  		s.Equal(expected, result)
  3154  	}
  3155  	s.Run("sanity check", func() {
  3156  		// workflow should run the child successfully normally...
  3157  		check(-1, false, "no err")
  3158  	})
  3159  	s.Run("canceled after child starts", func() {
  3160  		// ... and cancel the child when the child is canceled...
  3161  		check(30*time.Second, false, "canceled")
  3162  	})
  3163  	s.Run("canceled before child starts", func() {
  3164  		// ... and should not start the child (i.e. be canceled) when canceled before it is started.
  3165  		check(0, false, "canceled")
  3166  	})
  3167  	s.Run("canceled before child starts with bugport enabled", func() {
  3168  		// prior to v0.18.4, canceling before the child was started would still start the child,
  3169  		// and it would continue running.
  3170  		// the bugport provides this old behavior to ease migration, at least until we feel the need to remove it.
  3171  		check(0, true, "no err")
  3172  	})
  3173  }
  3174  
  3175  func TestRegression_LocalActivityErrorEncoding(t *testing.T) {
  3176  	// previously not encoded correctly
  3177  	s := WorkflowTestSuite{}
  3178  	s.SetLogger(zaptest.NewLogger(t))
  3179  	env := s.NewTestWorkflowEnvironment()
  3180  	sentinel := errors.New("sentinel error value")
  3181  	env.RegisterWorkflowWithOptions(func(ctx Context) error {
  3182  		ctx = WithLocalActivityOptions(ctx, LocalActivityOptions{ScheduleToCloseTimeout: time.Second})
  3183  		err := ExecuteLocalActivity(ctx, func(ctx context.Context) error {
  3184  			return sentinel
  3185  		}).Get(ctx, nil)
  3186  		if errors.Is(err, sentinel) {
  3187  			// incorrect path, taken through v0.19.1
  3188  			return fmt.Errorf("local activity errors need to be encoded, and must not be `.Is` a specific value: %w", err)
  3189  		}
  3190  		// correct path
  3191  		return sentinel
  3192  	}, RegisterWorkflowOptions{Name: "errorsis"})
  3193  	env.ExecuteWorkflow("errorsis")
  3194  	err := env.GetWorkflowError()
  3195  
  3196  	// make sure that the right path was chosen, and that the GetWorkflowError returns the same encoded value.
  3197  	// GetWorkflowError could... *possibly* return a wrapped original error value, but this seems unnecessarily
  3198  	// difficult to maintain long-term, doesn't reflect actual non-test behavior, and might lead to misunderstandings.
  3199  	require.Error(t, err) // stop early to avoid confusing NPEs
  3200  	var generr *GenericError
  3201  	assert.ErrorAs(t, err, &generr, "should be an encoded generic error")
  3202  	assert.NotErrorIs(t, err, sentinel, "should not contain a specific value, as this cannot be replayed")
  3203  	assert.Contains(t, err.Error(), "sentinel error value", "should contain the user error text")
  3204  	assert.NotContains(t, err.Error(), "need to be encoded", "should not contain the wrong-err-type branch message")
  3205  }