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

Code before:
    @Test
    public void shouldShutdownTaskExecutor() {
        assertNull(taskExecutor.currentTask(), "Have task assigned before startup");

        taskExecutor.start();

        verify(taskManager, timeout(VERIFICATION_TIMEOUT)).assignNextTask(taskExecutor);

        taskExecutor.shutdown(Duration.ofMinutes(1));

        verify(task).prepareCommit();
        verify(taskManager).unassignTask(task, taskExecutor);

        assertNull(taskExecutor.currentTask(), "Have task assigned after shutdown");
    }

    @Test
    public void shouldAwaitProcessableTasksIfNoneAssignable() throws InterruptedException {
        assertNull(taskExecutor.currentTask(), "Have task assigned before startup");
        when(task.isProcessable(anyLong())).thenReturn(false);

        taskExecutor.start();

        verify(taskManager, timeout(VERIFICATION_TIMEOUT).atLeastOnce()).awaitProcessableTasks();
    }

    @Test
    public void shouldUnassignTaskWhenNotProgressing() {
        when(task.isProcessable(anyLong())).thenReturn(false);
        when(task.maybePunctuateStreamTime()).thenReturn(false);
        when(task.maybePunctuateSystemTime()).thenReturn(false);

        taskExecutor.start();

        verify(taskManager, timeout(VERIFICATION_TIMEOUT)).unassignTask(task, taskExecutor);
        verify(task).prepareCommit();
        assertNull(taskExecutor.currentTask());
    }

    @Test
    public void shouldProcessTasks() {
        when(taskExecutionMetadata.canProcessTask(any(), anyLong())).thenReturn(true);
        when(task.isProcessable(anyLong())).thenReturn(true);

        taskExecutor.start();

        verify(task, timeout(VERIFICATION_TIMEOUT).atLeast(2)).process(anyLong());
        verify(task, timeout(VERIFICATION_TIMEOUT).atLeastOnce()).recordProcessBatchTime(anyLong());
    }



Review comment (suggestion):
This test is a bit hard to follow. I think the test would be easier to understand if you specify the simpler stub `when(taskManager.assignNextTask(taskExecutor)).thenReturn(null);` here in the test instead of using the stub specified in `setUp()`.   

Code after:
    @Test
    public void shouldShutdownTaskExecutor() {
        assertNull(taskExecutor.currentTask(), "Have task assigned before startup");

        taskExecutor.start();

        verify(taskManager, timeout(VERIFICATION_TIMEOUT)).assignNextTask(taskExecutor);

        taskExecutor.shutdown(Duration.ofMinutes(1));

        verify(task).prepareCommit();
        verify(taskManager).unassignTask(task, taskExecutor);

        assertNull(taskExecutor.currentTask(), "Have task assigned after shutdown");
    }

    @Test
    public void shouldAwaitProcessableTasksIfNoneAssignable() throws InterruptedException {
        assertNull(taskExecutor.currentTask(), "Have task assigned before startup");
        when(taskManager.assignNextTask(taskExecutor)).thenReturn(null);

        taskExecutor.start();

        verify(taskManager, timeout(VERIFICATION_TIMEOUT).atLeastOnce()).awaitProcessableTasks();
    }

    @Test
    public void shouldUnassignTaskWhenNotProgressing() {
        when(task.isProcessable(anyLong())).thenReturn(false);
        when(task.maybePunctuateStreamTime()).thenReturn(false);
        when(task.maybePunctuateSystemTime()).thenReturn(false);

        taskExecutor.start();

        verify(taskManager, timeout(VERIFICATION_TIMEOUT)).unassignTask(task, taskExecutor);
        verify(task).prepareCommit();
        assertNull(taskExecutor.currentTask());
    }

    @Test
    public void shouldProcessTasks() {
        when(taskExecutionMetadata.canProcessTask(any(), anyLong())).thenReturn(true);
        when(task.isProcessable(anyLong())).thenReturn(true);

        taskExecutor.start();

        verify(task, timeout(VERIFICATION_TIMEOUT).atLeast(2)).process(anyLong());
        verify(task, timeout(VERIFICATION_TIMEOUT).atLeastOnce()).recordProcessBatchTime(anyLong());
    }