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

Code before:
    }

    @Test
    public void shouldPropagateKafkaExceptionsInPunctuateActiveTasks() {
        final StateMachineTask task00 = new StateMachineTask(taskId00, taskId00Partitions, true, stateManager) {
            @Override
            public boolean maybePunctuateStreamTime() {
                throw new KafkaException("oops");
            }
        };

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

        taskManager.handleAssignment(taskId00Assignment, emptyMap());
        assertThat(taskManager.tryToCompleteRestoration(time.milliseconds(), null), is(true));

        assertThat(task00.state(), is(Task.State.RUNNING));

        assertThrows(KafkaException.class, () -> taskManager.punctuate());
    }

    @Test
    public void shouldPunctuateActiveTasks() {

        final StreamTask task00 = statefulTask(taskId00, taskId00ChangelogPartitions)
            .withInputPartitions(taskId00Partitions)
            .inState(State.RUNNING)
            .build();

        when(task00.maybePunctuateStreamTime()).thenReturn(true);
        when(task00.maybePunctuateSystemTime()).thenReturn(true);

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

        when(stateUpdater.restoresActiveTasks()).thenReturn(false);
        when(stateUpdater.hasExceptionsAndFailedTasks()).thenReturn(false);
        when(stateUpdater.drainRestoredActiveTasks(any(Duration.class))).thenReturn(Set.of());
        when(stateUpdater.drainExceptionsAndFailedTasks()).thenReturn(List.of());

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

        assertTrue(taskManager.checkStateUpdater(time.milliseconds(), noOpResetter));

        // one for stream and one for system time
        assertThat(taskManager.punctuate(), equalTo(2));

        verify(task00).maybePunctuateStreamTime();
        verify(task00).maybePunctuateSystemTime();
    }

Review comment (suggestion):
this test seems to be testing what it should. question is just whether it can be simplified (see below)

Code after:

        assertThrows(TaskMigratedException.class, () -> taskManager.punctuate());
    }

    @Test
    public void shouldPropagateKafkaExceptionsInPunctuateActiveTasks() {
        final StateMachineTask task00 = new StateMachineTask(taskId00, taskId00Partitions, true, stateManager) {
            @Override
            public boolean maybePunctuateStreamTime() {
                throw new KafkaException("oops");
            }
        };

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

        taskManager.handleAssignment(taskId00Assignment, emptyMap());
        assertThat(taskManager.tryToCompleteRestoration(time.milliseconds(), null), is(true));

        assertThat(task00.state(), is(Task.State.RUNNING));

        assertThrows(KafkaException.class, () -> taskManager.punctuate());
    }

    @Test
    public void shouldPunctuateActiveTasks() {

        final StreamTask task00 = statefulTask(taskId00, taskId00ChangelogPartitions)
            .withInputPartitions(taskId00Partitions)
            .inState(State.RUNNING)
            .build();

        when(task00.maybePunctuateStreamTime()).thenReturn(true);
        when(task00.maybePunctuateSystemTime()).thenReturn(true);

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

        when(stateUpdater.restoresActiveTasks()).thenReturn(false);
        when(stateUpdater.hasExceptionsAndFailedTasks()).thenReturn(false);
        when(stateUpdater.drainRestoredActiveTasks(any(Duration.class))).thenReturn(Set.of());
        when(stateUpdater.drainExceptionsAndFailedTasks()).thenReturn(List.of());

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

        // one for stream and one for system time
        assertThat(taskManager.punctuate(), equalTo(2));

        verify(task00).maybePunctuateStreamTime();
        verify(task00).maybePunctuateSystemTime();
    }