Fix schedule workflow unit test

temporalio · Oct 31, 2023 · 9c0c3b9 · 9c0c3b9
1 parent 2320fc1
commit 9c0c3b9
Showing 1 changed file with 9 additions and 0 deletions.
diff --git a/service/worker/scheduler/workflow_test.go b/service/worker/scheduler/workflow_test.go
@@ -1718,6 +1718,9 @@ func (s *workflowSuite) TestCANByIterations() {
 		s.Fail("too many starts")
 		return nil, nil
 	}).Times(0).Maybe()
+	s.expectWatch(func(req *schedspb.WatchWorkflowRequest) (*schedspb.WatchWorkflowResponse, error) {
+		return &schedspb.WatchWorkflowResponse{Status: enumspb.WORKFLOW_EXECUTION_STATUS_COMPLETED}, nil
+	}).Times(0).Maybe()
 
 	// this is ignored because we set iters explicitly
 	s.env.RegisterDelayedCallback(func() {
@@ -1755,6 +1758,9 @@ func (s *workflowSuite) TestCANBySuggested() {
 		s.Fail("too many starts", req.Request.WorkflowId)
 		return nil, nil
 	}).Times(0).Maybe()
+	s.expectWatch(func(req *schedspb.WatchWorkflowRequest) (*schedspb.WatchWorkflowResponse, error) {
+		return &schedspb.WatchWorkflowResponse{Status: enumspb.WORKFLOW_EXECUTION_STATUS_COMPLETED}, nil
+	}).Times(0).Maybe()
 
 	s.env.RegisterDelayedCallback(func() {
 		s.env.SetContinueAsNewSuggested(true)
@@ -1791,6 +1797,9 @@ func (s *workflowSuite) TestCANBySignal() {
 		s.Fail("too many starts", req.Request.WorkflowId)
 		return nil, nil
 	}).Times(0).Maybe()
+	s.expectWatch(func(req *schedspb.WatchWorkflowRequest) (*schedspb.WatchWorkflowResponse, error) {
+		return &schedspb.WatchWorkflowResponse{Status: enumspb.WORKFLOW_EXECUTION_STATUS_COMPLETED}, nil
+	}).Times(0).Maybe()
 
 	s.env.RegisterDelayedCallback(func() {
 		s.env.SignalWorkflow(SignalNameForceCAN, nil)