diff --git a/internal/internal_workflow_testsuite.go b/internal/internal_workflow_testsuite.go index e5104581f..8270a2217 100644 --- a/internal/internal_workflow_testsuite.go +++ b/internal/internal_workflow_testsuite.go @@ -169,6 +169,9 @@ type ( onTimerScheduledListener func(timerID string, duration time.Duration) onTimerFiredListener func(timerID string) onTimerCanceledListener func(timerID string) + + activityScheduleToCloseTimeout time.Duration + activityStartToCloseTimeout time.Duration } // testWorkflowEnvironmentImpl is the environment that runs the workflow/activity unit tests. @@ -236,6 +239,9 @@ func newTestWorkflowEnvironmentImpl(s *WorkflowTestSuite, parentRegistry *regist callbackChannel: make(chan testCallbackHandle, 1000), testTimeout: 3 * time.Second, expectedMockCalls: make(map[string]struct{}), + + activityScheduleToCloseTimeout: 600 * time.Second, + activityStartToCloseTimeout: 600 * time.Second, }, workflowInfo: &WorkflowInfo{ @@ -544,8 +550,8 @@ func (env *testWorkflowEnvironmentImpl) executeActivity( parameters := ExecuteActivityParams{ ExecuteActivityOptions: ExecuteActivityOptions{ - ScheduleToCloseTimeout: 600 * time.Second, - StartToCloseTimeout: 600 * time.Second, + ScheduleToCloseTimeout: env.activityScheduleToCloseTimeout, + StartToCloseTimeout: env.activityStartToCloseTimeout, }, ActivityType: *activityType, Input: input, diff --git a/internal/internal_workflow_testsuite_test.go b/internal/internal_workflow_testsuite_test.go index 916372d91..c5d3d98be 100644 --- a/internal/internal_workflow_testsuite_test.go +++ b/internal/internal_workflow_testsuite_test.go @@ -4031,3 +4031,49 @@ func (s *WorkflowTestSuiteUnitTest) Test_WorkflowGetCurrentHistoryLength() { s.NoError(env.GetWorkflowResult(&result)) s.Equal(17, result) } + +func (s *WorkflowTestSuiteUnitTest) Test_ActivityWithStartToCloseTimeout() { + timeout := 100 * time.Millisecond + + timeoutActivity := func(ctx context.Context) error { + time.Sleep(timeout * 2) + return nil + } + + noTimeoutActivity := func(ctx context.Context) error { + time.Sleep(timeout / 2) + return nil + } + + env := s.NewTestActivityEnvironment() + env.SetActivityStartToCloseTimeout(timeout) + env.RegisterActivity(timeoutActivity) + env.RegisterActivity(noTimeoutActivity) + _, err := env.ExecuteActivity(timeoutActivity) + s.Error(err) + _, err = env.ExecuteActivity(noTimeoutActivity) + s.NoError(err) +} + +func (s *WorkflowTestSuiteUnitTest) Test_ActivityWithScheduleToCloseTimeout() { + timeout := 100 * time.Millisecond + + timeoutActivity := func(ctx context.Context) error { + time.Sleep(timeout * 2) + return nil + } + + noTimeoutActivity := func(ctx context.Context) error { + time.Sleep(timeout / 2) + return nil + } + + env := s.NewTestActivityEnvironment() + env.SetActivityScheduleToCloseTimeout(timeout) + env.RegisterActivity(timeoutActivity) + env.RegisterActivity(noTimeoutActivity) + _, err := env.ExecuteActivity(timeoutActivity) + s.Error(err) + _, err = env.ExecuteActivity(noTimeoutActivity) + s.NoError(err) +} diff --git a/internal/workflow_testsuite.go b/internal/workflow_testsuite.go index 2b633f3f2..79f5b6941 100644 --- a/internal/workflow_testsuite.go +++ b/internal/workflow_testsuite.go @@ -262,6 +262,16 @@ func (t *TestActivityEnvironment) SetOnActivityHeartbeatListener( return t } +// SetActivityStartToCloseTimeout sets the start to close timeouts of activities in the test environment. +func (e *TestActivityEnvironment) SetActivityStartToCloseTimeout(timeout time.Duration) { + e.impl.testWorkflowEnvironmentShared.activityStartToCloseTimeout = timeout +} + +// SetActivityStartToCloseTimeout sets the schedule to close timeouts of activities in the test environment +func (e *TestActivityEnvironment) SetActivityScheduleToCloseTimeout(timeout time.Duration) { + e.impl.testWorkflowEnvironmentShared.activityScheduleToCloseTimeout = timeout +} + // RegisterWorkflow registers workflow implementation with the TestWorkflowEnvironment func (e *TestWorkflowEnvironment) RegisterWorkflow(w interface{}) { e.impl.RegisterWorkflow(w)