From 0d10acf0ad5aeefa2afed685fb846165996118d1 Mon Sep 17 00:00:00 2001
From: ohad83 <ohad83@gmail.com>
Date: Sat, 5 Aug 2023 22:35:27 +0300
Subject: [PATCH] Allow setting activity timeouts in testsuite

Add 2 method to TestActivityEnvironment:
- SetActivityStartToCloseTimeout
- SetActivityScheduleToCloseTimeout

This will allow changing the default 10 minute timeouts in the current test environments.
---
 internal/internal_workflow_testsuite.go      | 10 ++++-
 internal/internal_workflow_testsuite_test.go | 46 ++++++++++++++++++++
 internal/workflow_testsuite.go               | 10 +++++
 3 files changed, 64 insertions(+), 2 deletions(-)

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)