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

Code before:
        doThrow(taskCorruptedException).doNothing().when(changelogReader).restore(updatingTasks1);
        when(changelogReader.allChangelogsCompleted()).thenReturn(false);
        stateUpdater.start();

        stateUpdater.add(activeTask1);
        stateUpdater.add(activeTask2);
        stateUpdater.add(standbyTask);

        final ExceptionAndTasks expectedExceptionAndTasks =
            new ExceptionAndTasks(mkSet(activeTask1, activeTask2), taskCorruptedException);
        verifyExceptionsAndFailedTasks(expectedExceptionAndTasks);
        final InOrder orderVerifier = inOrder(changelogReader);
        orderVerifier.verify(changelogReader, atLeast(1)).enforceRestoreActive();
        orderVerifier.verify(changelogReader).transitToUpdateStandby();
    }
    
    @Test
    public void shouldNotTransitToStandbyAgainAfterStandbyTaskFailed() throws Exception {
        final StandbyTask task1 = standbyTask(TASK_0_0, mkSet(TOPIC_PARTITION_A_0)).inState(State.RUNNING).build();
        final StandbyTask task2 = standbyTask(TASK_1_0, mkSet(TOPIC_PARTITION_B_0)).inState(State.RUNNING).build();
        final Map<TaskId, Task> updatingTasks = mkMap(
                mkEntry(task1.id(), task1),
                mkEntry(task2.id(), task2)
        );
        final TaskCorruptedException taskCorruptedException =
                new TaskCorruptedException(mkSet(task1.id()));
        final ExceptionAndTasks expectedExceptionAndTasks =
                new ExceptionAndTasks(mkSet(task1), taskCorruptedException);
        when(changelogReader.allChangelogsCompleted()).thenReturn(false);
        doThrow(taskCorruptedException).doNothing().when(changelogReader).restore(updatingTasks);

        stateUpdater.start();
        stateUpdater.add(task1);
        stateUpdater.add(task2);

        verifyExceptionsAndFailedTasks(expectedExceptionAndTasks);
        verify(changelogReader, times(1)).transitToUpdateStandby();
    }

    @Test
    public void shouldUpdateStandbyTaskAfterAllActiveStatefulTasksRemoved() throws Exception {
        final StreamTask activeTask1 = statefulTask(TASK_0_0, mkSet(TOPIC_PARTITION_A_0)).inState(State.RESTORING).build();
        final StreamTask activeTask2 = statefulTask(TASK_0_1, mkSet(TOPIC_PARTITION_B_0)).inState(State.RESTORING).build();
        final StandbyTask standbyTask = standbyTask(TASK_1_0, mkSet(TOPIC_PARTITION_C_0)).inState(State.RUNNING).build();
        when(changelogReader.completedChangelogs()).thenReturn(Collections.emptySet());
        when(changelogReader.allChangelogsCompleted()).thenReturn(false);
        stateUpdater.start();
        stateUpdater.add(activeTask1);
        stateUpdater.add(activeTask2);
        stateUpdater.add(standbyTask);
        verifyUpdatingTasks(activeTask1, activeTask2, standbyTask);

Review comment (suggestion):
nit:
```suggestion
        final TaskCorruptedException taskCorruptedException = new TaskCorruptedException(mkSet(task1.id()));
```

Code after:
        doThrow(taskCorruptedException).doNothing().when(changelogReader).restore(updatingTasks1);
        when(changelogReader.allChangelogsCompleted()).thenReturn(false);
        stateUpdater.start();

        stateUpdater.add(activeTask1);
        stateUpdater.add(activeTask2);
        stateUpdater.add(standbyTask);

        final ExceptionAndTasks expectedExceptionAndTasks =
            new ExceptionAndTasks(mkSet(activeTask1, activeTask2), taskCorruptedException);
        verifyExceptionsAndFailedTasks(expectedExceptionAndTasks);
        final InOrder orderVerifier = inOrder(changelogReader);
        orderVerifier.verify(changelogReader, atLeast(1)).enforceRestoreActive();
        orderVerifier.verify(changelogReader).transitToUpdateStandby();
    }
    
    @Test
    public void shouldNotTransitToStandbyAgainAfterStandbyTaskFailed() throws Exception {
        final StandbyTask task1 = standbyTask(TASK_0_0, mkSet(TOPIC_PARTITION_A_0)).inState(State.RUNNING).build();
        final StandbyTask task2 = standbyTask(TASK_1_0, mkSet(TOPIC_PARTITION_B_0)).inState(State.RUNNING).build();
        final Map<TaskId, Task> updatingTasks = mkMap(
                mkEntry(task1.id(), task1),
                mkEntry(task2.id(), task2)
        );
        final TaskCorruptedException taskCorruptedException = new TaskCorruptedException(mkSet(task1.id()));
        final ExceptionAndTasks expectedExceptionAndTasks = new ExceptionAndTasks(mkSet(task1), taskCorruptedException);
        when(changelogReader.allChangelogsCompleted()).thenReturn(false);
        doThrow(taskCorruptedException).doNothing().when(changelogReader).restore(updatingTasks);

        stateUpdater.start();
        stateUpdater.add(task1);
        stateUpdater.add(task2);

        verifyExceptionsAndFailedTasks(expectedExceptionAndTasks);
        verify(changelogReader, times(1)).transitToUpdateStandby();
    }

    @Test
    public void shouldUpdateStandbyTaskAfterAllActiveStatefulTasksRemoved() throws Exception {
        final StreamTask activeTask1 = statefulTask(TASK_0_0, mkSet(TOPIC_PARTITION_A_0)).inState(State.RESTORING).build();
        final StreamTask activeTask2 = statefulTask(TASK_0_1, mkSet(TOPIC_PARTITION_B_0)).inState(State.RESTORING).build();
        final StandbyTask standbyTask = standbyTask(TASK_1_0, mkSet(TOPIC_PARTITION_C_0)).inState(State.RUNNING).build();
        when(changelogReader.completedChangelogs()).thenReturn(Collections.emptySet());
        when(changelogReader.allChangelogsCompleted()).thenReturn(false);
        stateUpdater.start();
        stateUpdater.add(activeTask1);
        stateUpdater.add(activeTask2);
        stateUpdater.add(standbyTask);
        verifyUpdatingTasks(activeTask1, activeTask2, standbyTask);

        stateUpdater.remove(activeTask1.id());