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

Code before:
                throw new RuntimeException("oops");
            }
        };

        when(consumer.assignment())
            .thenReturn(assignment)
            .thenReturn(taskId00Partitions);
        when(activeTaskCreator.createTasks(any(), eq(taskId00Assignment))).thenReturn(singletonList(task00));

        taskManager.handleAssignment(taskId00Assignment, emptyMap());
        assertThat(taskManager.tryToCompleteRestoration(time.milliseconds(), tp -> assertThat(tp, is(empty()))), is(true));
        assertThat(task00.state(), is(Task.State.RUNNING));

        task00.setChangelogOffsets(singletonMap(t1p0, 0L));
        taskManager.handleCorruption(singleton(taskId00));
        assertThat(task00.commitPrepared, is(true));
        assertThat(task00.state(), is(Task.State.CREATED));
        assertThat(task00.partitionsForOffsetReset, equalTo(taskId00Partitions));
        assertThat(taskManager.activeTaskMap(), is(singletonMap(taskId00, task00)));
        assertThat(taskManager.standbyTaskMap(), Matchers.anEmptyMap());

        verify(stateManager).markChangelogAsCorrupted(taskId00Partitions);
    }

    @Test
    public void shouldNotCommitCorruptedTasksOnTaskCorruptedException() {
        final StreamTask corruptedTask = statefulTask(taskId00, taskId00ChangelogPartitions)
            .withInputPartitions(taskId00Partitions)
            .inState(State.RUNNING)
            .build();

        final TaskCorruptedException corruptedException = new TaskCorruptedException(singleton(taskId00));
        final ExceptionAndTask exceptionAndTask = new ExceptionAndTask(corruptedException, corruptedTask);

        when(stateUpdater.hasExceptionsAndFailedTasks()).thenReturn(true);
        when(stateUpdater.drainExceptionsAndFailedTasks()).thenReturn(singletonList(exceptionAndTask));

        final TasksRegistry tasks = mock(TasksRegistry.class);

        final TaskManager taskManager = setUpTaskManager(ProcessingMode.AT_LEAST_ONCE, tasks, false);

        final TaskCorruptedException thrown = assertThrows(
            TaskCorruptedException.class,
            () -> taskManager.checkStateUpdater(time.milliseconds(), noOpResetter)
        );

        assertEquals(singleton(taskId00), thrown.corruptedTasks());
        assertEquals("Tasks [0_0] are corrupted and hence need to be re-initialized", thrown.getMessage());

        verify(tasks).addFailedTask(corruptedTask);


Review comment (bug):
I renamed this test from `shouldCommitNonCorruptedTasksOnTaskCorruptedException`. Based on my understanding, the commit logic happens at the `StreamThread` level, but only the exception propagation happens in `TaskManager` with `checkStateUpdater`. So I decided to omit the check for commit logic and rewrite the test. 

And hence I propose to rename to `shouldNotCommitCorruptedTasksOnTaskCorruptedException`

Please correct if I am wrong or If I misunderstood!


Code after:
        when(activeTaskCreator.createTasks(any(), eq(taskId00Assignment))).thenReturn(singletonList(task00));

        taskManager.handleAssignment(taskId00Assignment, emptyMap());
        assertThat(taskManager.tryToCompleteRestoration(time.milliseconds(), tp -> assertThat(tp, is(empty()))), is(true));
        assertThat(task00.state(), is(Task.State.RUNNING));

        task00.setChangelogOffsets(singletonMap(t1p0, 0L));
        taskManager.handleCorruption(singleton(taskId00));
        assertThat(task00.commitPrepared, is(true));
        assertThat(task00.state(), is(Task.State.CREATED));
        assertThat(task00.partitionsForOffsetReset, equalTo(taskId00Partitions));
        assertThat(taskManager.activeTaskMap(), is(singletonMap(taskId00, task00)));
        assertThat(taskManager.standbyTaskMap(), Matchers.anEmptyMap());

        verify(stateManager).markChangelogAsCorrupted(taskId00Partitions);
    }

    @Test
    public void shouldCommitNonCorruptedTasksOnTaskCorruptedException() {
        final StreamTask corruptedTask = statefulTask(taskId00, taskId00ChangelogPartitions)
            .withInputPartitions(taskId00Partitions)
            .inState(State.RUNNING)
            .build();

        final StreamTask nonCorruptedTask = statefulTask(taskId01, taskId01ChangelogPartitions)
            .withInputPartitions(taskId01Partitions)
            .inState(State.RUNNING)
            .build();

        final TasksRegistry tasks = mock(TasksRegistry.class);
        when(tasks.task(taskId00)).thenReturn(corruptedTask);
        when(tasks.allTasksPerId()).thenReturn(mkMap(
            mkEntry(taskId00, corruptedTask),
            mkEntry(taskId01, nonCorruptedTask)
        ));
        when(tasks.activeTaskIds()).thenReturn(Set.of(taskId00, taskId01));

        when(nonCorruptedTask.commitNeeded()).thenReturn(true);
        when(nonCorruptedTask.prepareCommit(true)).thenReturn(emptyMap());
        when(corruptedTask.prepareCommit(false)).thenReturn(emptyMap());
        doNothing().when(corruptedTask).postCommit(anyBoolean());

        when(consumer.assignment()).thenReturn(taskId00Partitions);

        final TaskManager taskManager = setUpTaskManagerWithStateUpdater(ProcessingMode.AT_LEAST_ONCE, tasks);

        taskManager.handleCorruption(Set.of(taskId00));

        verify(nonCorruptedTask).prepareCommit(true);
        verify(nonCorruptedTask, never()).addPartitionsForOffsetReset(any());
        verify(corruptedTask).addPartitionsForOffsetReset(taskId00Partitions);