Set activity handle in testWorkflowEnvironment (#274)

temporalio · Oct 22, 2020 · c49e036 · c49e036
1 parent b8d2ed4
commit c49e036
Show file tree

Hide file tree

Showing 3 changed files with 15 additions and 0 deletions.
diff --git a/go.sum b/go.sum
@@ -144,6 +144,7 @@ golang.org/x/sys v0.0.0-20200323222414-85ca7c5b95cd/go.mod h1:h1NjWce9XRLGQEsW7w
 golang.org/x/sys v0.0.0-20200923182605-d9f96fdee20d h1:L/IKR6COd7ubZrs2oTnTi73IhgqJ71c9s80WsQnh0Es=
 golang.org/x/sys v0.0.0-20200923182605-d9f96fdee20d/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
 golang.org/x/sys v0.0.0-20200930185726-fdedc70b468f/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
+golang.org/x/sys v0.0.0-20201016160150-f659759dc4ca h1:mLWBs1i4Qi5cHWGEtn2jieJQ2qtwV/gT0A2zLrmzaoE=
 golang.org/x/sys v0.0.0-20201016160150-f659759dc4ca/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
 golang.org/x/text v0.3.0 h1:g61tztE5qeGQ89tm6NTjjM9VPIm088od1l6aSorWRWg=
 golang.org/x/text v0.3.0/go.mod h1:NqM8EUOU14njkJ3fqMW+pc6Ldnwhi/IjpwHt7yyuwOQ=

diff --git a/internal/internal_workflow_testsuite.go b/internal/internal_workflow_testsuite.go
@@ -558,6 +558,9 @@ func (env *testWorkflowEnvironmentImpl) executeActivity(
 
 	// ensure activityFn is registered to defaultTestTaskQueue
 	taskHandler := env.newTestActivityTaskHandler(defaultTestTaskQueue, env.GetDataConverter())
+	activityHandle := &testActivityHandle{callback: func(result *commonpb.Payloads, err error) {}, activityType: parameters.ActivityType.Name}
+	env.setActivityHandle(scheduleTaskAttr.GetActivityId(), activityHandle)
+
 	result, err := taskHandler.Execute(defaultTestTaskQueue, task)
 	if err != nil {
 		if err == context.DeadlineExceeded {

diff --git a/internal/internal_workflow_testsuite_test.go b/internal/internal_workflow_testsuite_test.go
@@ -3007,6 +3007,17 @@ func (s *WorkflowTestSuiteUnitTest) Test_ActivityWithProgress() {
 	s.Equal(lastProgress+1, newProgress)
 }
 
+func (s *WorkflowTestSuiteUnitTest) Test_ActivityHeartbeat() {
+	env := s.NewTestActivityEnvironment()
+	env.RegisterActivity(testActivityHeartbeat)
+	v, err := env.ExecuteActivity(testActivityHeartbeat, "result", 1*time.Millisecond)
+	s.NoError(err)
+	var result string
+	err = v.Get(&result)
+	s.NoError(err)
+	s.Equal("heartbeat_result", result)
+}
+
 func (s *WorkflowTestSuiteUnitTest) Test_ActivityGoexit() {
 	fn := func(ctx context.Context) error {
 		runtime.Goexit() // usually this is called by t.FailNow(), but can't call FailNow here since that would mark the test as failed.