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

Code before:
        doThrow(new RuntimeException("KABOOM!")).when(task00).suspend();

        final TasksRegistry tasks = mock(TasksRegistry.class);
        when(tasks.allTasks()).thenReturn(Set.of(task00));

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

        assertThrows(RuntimeException.class, () -> taskManager.handleRevocation(taskId00Partitions));

        verify(task00).suspend();
    }

    @Test
    public void shouldCloseActiveTasksAndPropagateExceptionsOnCleanShutdownWithAlos() {
        shouldCloseActiveTasksAndPropagateExceptionsOnCleanShutdown(ProcessingMode.AT_LEAST_ONCE);
    }

    @Test
    public void shouldCloseActiveTasksAndPropagateExceptionsOnCleanShutdownWithExactlyOnceV2() {
        when(activeTaskCreator.streamsProducer()).thenReturn(mock(StreamsProducer.class));
        shouldCloseActiveTasksAndPropagateExceptionsOnCleanShutdown(ProcessingMode.EXACTLY_ONCE_V2);
    }

    private void shouldCloseActiveTasksAndPropagateExceptionsOnCleanShutdown(final ProcessingMode processingMode) {

        final StreamTask task00 = statefulTask(taskId00, taskId00ChangelogPartitions)
            .inState(State.RESTORING)
            .withInputPartitions(taskId00Partitions).build();
        final StreamTask task01 = statefulTask(taskId01, taskId01ChangelogPartitions)
            .inState(State.RUNNING)
            .withInputPartitions(taskId01Partitions).build();
        final StreamTask task02 = statefulTask(taskId02, taskId02ChangelogPartitions)
            .inState(State.RUNNING)
            .withInputPartitions(taskId02Partitions).build();
        final StreamTask task03 = statefulTask(taskId03, taskId03ChangelogPartitions)
            .inState(State.RUNNING)
            .withInputPartitions(taskId03Partitions).build();

        final TasksRegistry tasks = mock(TasksRegistry.class);
        final TaskManager taskManager = setUpTaskManagerWithStateUpdater(processingMode, tasks);

        doThrow(new TaskMigratedException("migrated", new RuntimeException("cause")))
            .when(task01).suspend();
        doThrow(new RuntimeException("oops"))
            .when(task02).suspend();
        doThrow(new RuntimeException("oops"))
            .when(task03).suspend();

        when(tasks.activeTasks()).thenReturn(Set.of(task00, task01, task02, task03));

        final RuntimeException exception = assertThrows(

Review comment (suggestion):
This test is a bit ewird. IIUC, we have an active restoring task that is not owned by the state updater in this test. But this should be impossible. I think we shouldn't have `task00` like this.

Code after:
        doThrow(new RuntimeException("KABOOM!")).when(task00).suspend();

        final TasksRegistry tasks = mock(TasksRegistry.class);
        when(tasks.allTasks()).thenReturn(Set.of(task00));

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

        assertThrows(RuntimeException.class, () -> taskManager.handleRevocation(taskId00Partitions));

        verify(task00).suspend();
    }

    @Test
    public void shouldCloseActiveTasksAndPropagateExceptionsOnCleanShutdownWithAlos() {
        shouldCloseActiveTasksAndPropagateExceptionsOnCleanShutdown(ProcessingMode.AT_LEAST_ONCE);
    }

    @Test
    public void shouldCloseActiveTasksAndPropagateExceptionsOnCleanShutdownWithExactlyOnceV2() {
        when(activeTaskCreator.streamsProducer()).thenReturn(mock(StreamsProducer.class));
        shouldCloseActiveTasksAndPropagateExceptionsOnCleanShutdown(ProcessingMode.EXACTLY_ONCE_V2);
    }

    private void shouldCloseActiveTasksAndPropagateExceptionsOnCleanShutdown(final ProcessingMode processingMode) {

        final StreamTask task00 = statefulTask(taskId00, taskId00ChangelogPartitions)
            .inState(State.RUNNING)
            .withInputPartitions(taskId00Partitions).build();
        final StreamTask task01 = statefulTask(taskId01, taskId01ChangelogPartitions)
            .inState(State.RUNNING)
            .withInputPartitions(taskId01Partitions).build();
        final StreamTask task02 = statefulTask(taskId02, taskId02ChangelogPartitions)
            .inState(State.RUNNING)
            .withInputPartitions(taskId02Partitions).build();

        final TasksRegistry tasks = mock(TasksRegistry.class);
        final TaskManager taskManager = setUpTaskManagerWithStateUpdater(processingMode, tasks);

        doThrow(new TaskMigratedException("migrated", new RuntimeException("cause")))
            .when(task01).suspend();
        doThrow(new RuntimeException("oops"))
            .when(task02).suspend();

        when(tasks.activeTasks()).thenReturn(Set.of(task00, task01, task02));

        final RuntimeException exception = assertThrows(
            RuntimeException.class,
            () -> taskManager.shutdown(true)
        );
        assertThat(exception.getCause().getMessage(), is("oops"));