Repository: apache/kafka
File: streams/src/test/java/org/apache/kafka/streams/processor/internals/TaskManagerTest.java (Java)

Code before:
        taskManager.checkStateUpdater(time.milliseconds(), noOpResetter);

        verify(consumer);
        Mockito.verify(statefulTask).suspend();
        Mockito.verify(tasks).addTask(statefulTask);
    }

    @Test
    public void shouldHandleMultipleRemovedTasksFromStateUpdater() {
        final StreamTask taskToRecycle0 = statefulTask(taskId00, taskId00ChangelogPartitions)
            .inState(State.RESTORING)
            .withInputPartitions(taskId00Partitions).build();
        final StandbyTask taskToRecycle1 = standbyTask(taskId01, taskId01ChangelogPartitions)
            .inState(State.RUNNING)
            .withInputPartitions(taskId01Partitions).build();
        final StandbyTask convertedTask0 = standbyTask(taskId00, taskId00ChangelogPartitions).build();
        final StreamTask convertedTask1 = statefulTask(taskId01, taskId01ChangelogPartitions).build();
        final StreamTask taskToClose = statefulTask(taskId02, taskId02ChangelogPartitions)
            .inState(State.RESTORING)
            .withInputPartitions(taskId02Partitions).build();
        final StreamTask taskToUpdateInputPartitions = statefulTask(taskId03, taskId03ChangelogPartitions)
            .inState(State.RESTORING)
            .withInputPartitions(taskId03Partitions).build();
        final StreamTask taskToCloseReviveAndUpdateInputPartitions = statefulTask(taskId04, taskId04ChangelogPartitions)
                .inState(State.RESTORING)
                .withInputPartitions(taskId04Partitions).build();
        when(stateUpdater.hasRemovedTasks()).thenReturn(true);
        when(stateUpdater.drainRemovedTasks())
            .thenReturn(mkSet(taskToRecycle0, taskToRecycle1, taskToClose, taskToUpdateInputPartitions, taskToCloseReviveAndUpdateInputPartitions));
        when(stateUpdater.restoresActiveTasks()).thenReturn(true);
        when(activeTaskCreator.createActiveTaskFromStandby(taskToRecycle1, taskId01Partitions, consumer))
            .thenReturn(convertedTask1);
        when(standbyTaskCreator.createStandbyTaskFromActive(taskToRecycle0, taskId00Partitions))
            .thenReturn(convertedTask0);
        expect(consumer.assignment()).andReturn(emptySet()).anyTimes();
        consumer.resume(anyObject());
        expectLastCall().anyTimes();
        final TasksRegistry tasks = mock(TasksRegistry.class);
        when(tasks.removePendingTaskToCloseClean(taskToClose.id())).thenReturn(true);
        when(tasks.removePendingTaskToCloseClean(argThat(taskId -> !taskId.equals(taskToClose.id())))).thenReturn(false);
        when(tasks.removePendingTaskToRecycle(taskToRecycle0.id())).thenReturn(taskId00Partitions);
        when(tasks.removePendingTaskToRecycle(taskToRecycle1.id())).thenReturn(taskId01Partitions);
        when(tasks.removePendingTaskToRecycle(
            argThat(taskId -> !taskId.equals(taskToRecycle0.id()) && !taskId.equals(taskToRecycle1.id())))
        ).thenReturn(null);
        when(tasks.removePendingTaskToUpdateInputPartitions(taskToUpdateInputPartitions.id())).thenReturn(taskId04Partitions);
        when(tasks.removePendingTaskToCloseReviveAndUpdateInputPartitions(taskToCloseReviveAndUpdateInputPartitions.id())).thenReturn(taskId05Partitions);
        when(tasks.removePendingTaskToCloseReviveAndUpdateInputPartitions(
                argThat(taskId -> !taskId.equals(taskToCloseReviveAndUpdateInputPartitions.id()))
        )).thenReturn(null);
        final TaskManager taskManager = setUpTaskManager(ProcessingMode.AT_LEAST_ONCE, tasks, true);

Review comment (suggestion):
nit:
```suggestion
            .inState(State.RESTORING)
            .withInputPartitions(taskId04Partitions).build();
```

Code after:
    public void shouldHandleMultipleRemovedTasksFromStateUpdater() {
        final StreamTask taskToRecycle0 = statefulTask(taskId00, taskId00ChangelogPartitions)
            .inState(State.RESTORING)
            .withInputPartitions(taskId00Partitions).build();
        final StandbyTask taskToRecycle1 = standbyTask(taskId01, taskId01ChangelogPartitions)
            .inState(State.RUNNING)
            .withInputPartitions(taskId01Partitions).build();
        final StandbyTask convertedTask0 = standbyTask(taskId00, taskId00ChangelogPartitions).build();
        final StreamTask convertedTask1 = statefulTask(taskId01, taskId01ChangelogPartitions).build();
        final StreamTask taskToClose = statefulTask(taskId02, taskId02ChangelogPartitions)
            .inState(State.RESTORING)
            .withInputPartitions(taskId02Partitions).build();
        final StreamTask taskToUpdateInputPartitions = statefulTask(taskId03, taskId03ChangelogPartitions)
            .inState(State.RESTORING)
            .withInputPartitions(taskId03Partitions).build();
        final StreamTask taskToCloseReviveAndUpdateInputPartitions = statefulTask(taskId04, taskId04ChangelogPartitions)
            .inState(State.RESTORING)
            .withInputPartitions(taskId04Partitions).build();
        when(stateUpdater.hasRemovedTasks()).thenReturn(true);
        when(stateUpdater.drainRemovedTasks())
            .thenReturn(mkSet(taskToRecycle0, taskToRecycle1, taskToClose, taskToUpdateInputPartitions, taskToCloseReviveAndUpdateInputPartitions));
        when(stateUpdater.restoresActiveTasks()).thenReturn(true);
        when(activeTaskCreator.createActiveTaskFromStandby(taskToRecycle1, taskId01Partitions, consumer))
            .thenReturn(convertedTask1);
        when(standbyTaskCreator.createStandbyTaskFromActive(taskToRecycle0, taskId00Partitions))
            .thenReturn(convertedTask0);
        expect(consumer.assignment()).andReturn(emptySet()).anyTimes();
        consumer.resume(anyObject());
        expectLastCall().anyTimes();
        final TasksRegistry tasks = mock(TasksRegistry.class);
        when(tasks.removePendingTaskToCloseClean(taskToClose.id())).thenReturn(true);
        when(tasks.removePendingTaskToCloseClean(argThat(taskId -> !taskId.equals(taskToClose.id())))).thenReturn(false);
        when(tasks.removePendingTaskToRecycle(taskToRecycle0.id())).thenReturn(taskId00Partitions);
        when(tasks.removePendingTaskToRecycle(taskToRecycle1.id())).thenReturn(taskId01Partitions);
        when(tasks.removePendingTaskToRecycle(
            argThat(taskId -> !taskId.equals(taskToRecycle0.id()) && !taskId.equals(taskToRecycle1.id())))
        ).thenReturn(null);
        when(tasks.removePendingTaskToUpdateInputPartitions(taskToUpdateInputPartitions.id())).thenReturn(taskId04Partitions);
        when(tasks.removePendingTaskToCloseReviveAndUpdateInputPartitions(taskToCloseReviveAndUpdateInputPartitions.id())).thenReturn(taskId05Partitions);
        when(tasks.removePendingTaskToCloseReviveAndUpdateInputPartitions(
            argThat(taskId -> !taskId.equals(taskToCloseReviveAndUpdateInputPartitions.id()))
        )).thenReturn(null);
        final TaskManager taskManager = setUpTaskManager(ProcessingMode.AT_LEAST_ONCE, tasks, true);
        taskManager.setMainConsumer(consumer);
        replay(consumer);

        taskManager.checkStateUpdater(time.milliseconds(), noOpResetter -> { });

        verify(consumer);
        Mockito.verify(activeTaskCreator, times(3)).closeAndRemoveTaskProducerIfNeeded(any());
        Mockito.verify(convertedTask0).initializeIfNeeded();