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

Code before:
        when(task.canPunctuateSystemTime()).thenReturn(true);

        assertEquals(task, taskManager.assignNextTask(taskExecutor));
        assertNull(taskManager.assignNextTask(taskExecutor));
    }

    @Test
    public void shouldAssignTasksThatCanBeStreamTimePunctuated() {
        taskManager.add(Collections.singleton(task));
        when(tasks.activeTasks()).thenReturn(Collections.singleton(task));
        when(taskExecutionMetadata.canPunctuateTask(eq(task))).thenReturn(true);
        when(task.canPunctuateStreamTime()).thenReturn(true);

        assertEquals(task, taskManager.assignNextTask(taskExecutor));
        assertNull(taskManager.assignNextTask(taskExecutor));
    }

    @Test
    public void shouldNotAssignTasksForPunctuationIfPunctuationDisabled() {
        taskManager.add(Collections.singleton(task));
        when(tasks.activeTasks()).thenReturn(Collections.singleton(task));
        when(taskExecutionMetadata.canPunctuateTask(eq(task))).thenReturn(false);
        when(task.canPunctuateStreamTime()).thenReturn(true);
        when(task.canPunctuateSystemTime()).thenReturn(true);

        assertNull(taskManager.assignNextTask(taskExecutor));
    }

    @Test
    public void shouldNotAssignTasksForProcessingIfProcessingDisabled() {
        taskManager.add(Collections.singleton(task));
        when(tasks.activeTasks()).thenReturn(Collections.singleton(task));
        when(taskExecutionMetadata.canProcessTask(eq(task), anyLong())).thenReturn(false);
        when(task.isProcessable(anyLong())).thenReturn(true);

        assertNull(taskManager.assignNextTask(taskExecutor));
    }

    @Test
    public void shouldUnassignTask() {
        taskManager.add(Collections.singleton(task));
        when(tasks.activeTasks()).thenReturn(Collections.singleton(task));
        when(taskExecutionMetadata.canProcessTask(eq(task), anyLong())).thenReturn(true);

        assertEquals(task, taskManager.assignNextTask(taskExecutor));

        taskManager.unassignTask(task, taskExecutor);
        assertEquals(task, taskManager.assignNextTask(taskExecutor));
    }

    @Test

Review comment (refactor):
Could you please extract these two lines to a method named `ensureTaskMakesProgress()` or similar and use that method where progress is needed?
I was a bit puzzled when I saw the two calls. 

Code after:
        taskManager.add(Collections.singleton(task));
        when(tasks.activeTasks()).thenReturn(Collections.singleton(task));
        when(taskExecutionMetadata.canProcessTask(eq(task), anyLong())).thenReturn(true);
        when(task.canPunctuateSystemTime()).thenReturn(true);

        assertEquals(task, taskManager.assignNextTask(taskExecutor));
        assertNull(taskManager.assignNextTask(taskExecutor));
    }

    @Test
    public void shouldAssignTasksThatCanBeStreamTimePunctuated() {
        taskManager.add(Collections.singleton(task));
        when(tasks.activeTasks()).thenReturn(Collections.singleton(task));
        when(taskExecutionMetadata.canPunctuateTask(eq(task))).thenReturn(true);
        when(task.canPunctuateStreamTime()).thenReturn(true);

        assertEquals(task, taskManager.assignNextTask(taskExecutor));
        assertNull(taskManager.assignNextTask(taskExecutor));
    }

    @Test
    public void shouldNotAssignTasksIfUncaughtExceptionPresent() {
        taskManager.add(Collections.singleton(task));
        when(tasks.activeTasks()).thenReturn(Collections.singleton(task));
        ensureTaskMakesProgress();
        taskManager.assignNextTask(taskExecutor);
        taskManager.setUncaughtException(new StreamsException("Exception"), taskId);
        taskManager.unassignTask(task, taskExecutor);

        assertNull(taskManager.assignNextTask(taskExecutor));
    }

    @Test
    public void shouldNotAssignTasksForPunctuationIfPunctuationDisabled() {
        taskManager.add(Collections.singleton(task));
        when(tasks.activeTasks()).thenReturn(Collections.singleton(task));
        when(taskExecutionMetadata.canPunctuateTask(eq(task))).thenReturn(false);
        when(task.canPunctuateStreamTime()).thenReturn(true);
        when(task.canPunctuateSystemTime()).thenReturn(true);

        assertNull(taskManager.assignNextTask(taskExecutor));
    }

    @Test
    public void shouldNotAssignTasksForProcessingIfProcessingDisabled() {
        taskManager.add(Collections.singleton(task));
        when(tasks.activeTasks()).thenReturn(Collections.singleton(task));
        when(taskExecutionMetadata.canProcessTask(eq(task), anyLong())).thenReturn(false);
        when(task.isProcessable(anyLong())).thenReturn(true);

        assertNull(taskManager.assignNextTask(taskExecutor));