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

Code before:
public class DefaultTaskManagerTest {

    private final static long VERIFICATION_TIMEOUT = 15000;

    private final Time time = new MockTime(1L);
    private final StreamTask task = mock(StreamTask.class);
    private final TasksRegistry tasks = mock(TasksRegistry.class);
    private final TaskExecutor taskExecutor = mock(TaskExecutor.class);
    private final StreamsException exception = mock(StreamsException.class);
    private final TaskExecutionMetadata taskExecutionMetadata = mock(TaskExecutionMetadata.class);

    private final StreamsConfig config = new StreamsConfig(configProps());
    private final TaskManager taskManager = new DefaultTaskManager(time, "TaskManager", tasks, config,
        (taskManager, name, time, taskExecutionMetadata) -> taskExecutor, taskExecutionMetadata);

    private Properties configProps() {
        return mkObjectProperties(mkMap(
            mkEntry(StreamsConfig.APPLICATION_ID_CONFIG, "appId"),
            mkEntry(StreamsConfig.BOOTSTRAP_SERVERS_CONFIG, "localhost:2171"),
            mkEntry(StreamsConfig.PROCESSING_GUARANTEE_CONFIG, StreamsConfig.EXACTLY_ONCE_V2)
        ));
    }

    @BeforeEach
    public void setUp() {
        when(task.id()).thenReturn(new TaskId(0, 0, "A"));
        when(task.isProcessable(anyLong())).thenReturn(true);
        when(task.isActive()).thenReturn(true);
        when(tasks.activeTasks()).thenReturn(Collections.singleton(task));
        when(tasks.task(new TaskId(0, 0, "A"))).thenReturn(task);
    }

    @Test
    public void shouldAddTask() {
        taskManager.add(Collections.singleton(task));

        verify(tasks).addTask(task);
        when(tasks.activeTasks()).thenReturn(Collections.singleton(task));
        assertEquals(1, taskManager.getTasks().size());
    }

    @Test
    public void shouldAssignTaskThatCanBeProcessed() {
        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));
        assertNull(taskManager.assignNextTask(taskExecutor));
    }


Review comment (suggestion):
nit: 
You could reuse `org.apache.kafka.test.StreamsTestUtils.TaskBuilder.statelessTask` here instead of creating your own stubs. 

Code after:

public class DefaultTaskManagerTest {

    private final static long VERIFICATION_TIMEOUT = 15000;

    private final Time time = new MockTime(1L);
    private final TaskId taskId = new TaskId(0, 0, "A");
    private final StreamTask task = TaskBuilder.statelessTask(taskId).build();
    private final TasksRegistry tasks = mock(TasksRegistry.class);
    private final TaskExecutor taskExecutor = mock(TaskExecutor.class);
    private final StreamsException exception = mock(StreamsException.class);
    private final TaskExecutionMetadata taskExecutionMetadata = mock(TaskExecutionMetadata.class);

    private final StreamsConfig config = new StreamsConfig(configProps());
    private final TaskManager taskManager = new DefaultTaskManager(time, "TaskManager", tasks, config,
        (taskManager, name, time, taskExecutionMetadata) -> taskExecutor, taskExecutionMetadata);

    private Properties configProps() {
        return mkObjectProperties(mkMap(
            mkEntry(StreamsConfig.APPLICATION_ID_CONFIG, "appId"),
            mkEntry(StreamsConfig.BOOTSTRAP_SERVERS_CONFIG, "localhost:2171"),
            mkEntry(StreamsConfig.PROCESSING_GUARANTEE_CONFIG, StreamsConfig.EXACTLY_ONCE_V2)
        ));
    }

    @BeforeEach
    public void setUp() {
        when(task.isProcessable(anyLong())).thenReturn(true);
        when(tasks.activeTasks()).thenReturn(Collections.singleton(task));
        when(tasks.task(taskId)).thenReturn(task);
    }

    @Test
    public void shouldAddTask() {
        taskManager.add(Collections.singleton(task));

        verify(tasks).addTask(task);
        when(tasks.activeTasks()).thenReturn(Collections.singleton(task));
        assertEquals(1, taskManager.getTasks().size());
    }

    @Test
    public void shouldAssignTaskThatCanBeProcessed() {
        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));
        assertNull(taskManager.assignNextTask(taskExecutor));
    }