@SuppressWarnings("rawtypes")
@Test
public void testMultipleTasksExecution() {

    String outputKey = "KEY";
    Task task = testTask();
    Worker worker = mock(Worker.class);
    when(worker.getPollingInterval()).thenReturn(3000);
    when(worker.getTaskDefName()).thenReturn(TEST_TASK_DEF_NAME);
    when(worker.execute(any())).thenAnswer(new Answer() {
        private int count = 0;
        Map<String, Object> outputMap = new HashMap<>();
        public TaskResult answer(InvocationOnMock invocation) {
            // Sleep for 2 seconds to simulate task execution
            Uninterruptibles.sleepUninterruptibly(2, TimeUnit.SECONDS);
            TaskResult taskResult = new TaskResult(task);
            outputMap.put(outputKey, count++);
            taskResult.setOutputData(outputMap);
            return taskResult;
        }
    });
    TaskClient taskClient = Mockito.mock(TaskClient.class);
    TaskPollExecutor taskPollExecutor = new TaskPollExecutor(null, taskClient, 1, 1, new HashMap<>(), "test-worker-");
    when(taskClient.pollTask(any(), any(), any())).thenReturn(task);
    when(taskClient.ack(any(), any())).thenReturn(true);
    CountDownLatch latch = new CountDownLatch(3);
    doAnswer(new Answer() {
                 private int count = 0;
                 public TaskResult answer(InvocationOnMock invocation) {
                     Object[] args = invocation.getArguments();
                     TaskResult result = (TaskResult) args[0];
                     assertEquals(IN_PROGRESS, result.getStatus());
                     assertEquals(count, result.getOutputData().get(outputKey));
                     count++;
                     latch.countDown();
                     return null;
                 }
             }
    ).when(taskClient).updateTask(any());
    Executors.newSingleThreadScheduledExecutor()
        .scheduleAtFixedRate(() -> taskPollExecutor.pollAndExecute(worker), 0, 1, TimeUnit.SECONDS);
    Uninterruptibles.awaitUninterruptibly(latch);
    // execute() is called 3 times on the worker (once for each task)
    verify(worker, times(3)).execute(any());
    verify(taskClient, times(3)).updateTask(any());
}