From 732d43a34d1e85fa040d57cb8fc32bdc0b3e05c6 Mon Sep 17 00:00:00 2001 From: Yichao Yang Date: Sat, 4 Jun 2022 11:55:10 -0700 Subject: [PATCH] Fix task parallel processor test (#2951) --- common/tasks/parallel_processor_test.go | 10 ++++++++-- 1 file changed, 8 insertions(+), 2 deletions(-) diff --git a/common/tasks/parallel_processor_test.go b/common/tasks/parallel_processor_test.go index a767c8141c2..8b8916d1479 100644 --- 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)