Fix task parallel processor test (#2951)

temporalio · Jun 4, 2022 · 732d43a · 732d43a
1 parent fa2e06c
commit 732d43a
Showing 1 changed file with 8 additions and 2 deletions.
diff --git a/common/tasks/parallel_processor_test.go b/common/tasks/parallel_processor_test.go
@@ -103,14 +103,20 @@ func (s *parallelProcessorSuite) TestSubmitProcess_Running_FailExecution() {
 	testWaitGroup.Wait()
 }
 
-func (s *parallelProcessorSuite) TestSubmitProcess_Stopped_FailSubmission() {
+func (s *parallelProcessorSuite) TestSubmitProcess_Stopped_Submission() {
 	testWaitGroup := sync.WaitGroup{}
 	testWaitGroup.Add(1)
 
 	s.processor.Stop()
 
-	// drain immediately
 	mockTask := NewMockTask(s.controller)
+
+	// if task get picked up before worker goroutine receives the shutdown notification
+	mockTask.EXPECT().RetryPolicy().Return(s.retryPolicy).MaxTimes(1)
+	mockTask.EXPECT().Execute().Return(nil).MaxTimes(1)
+	mockTask.EXPECT().Ack().Do(func() { testWaitGroup.Done() }).MaxTimes(1)
+
+	// if task get drained
 	mockTask.EXPECT().Reschedule().Do(func() { testWaitGroup.Done() }).MaxTimes(1)
 
 	s.processor.Submit(mockTask)