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

Code before:
        taskManager.assignNextTask(taskExecutor);
        taskManager.setUncaughtException(exception, task.id());

        assertThrows(IllegalArgumentException.class, () -> taskManager.setUncaughtException(exception, task.id()));
    }

    @Test
    public void shouldReturnExceptionsOnDrainExceptions() {
        taskManager.add(Collections.singleton(task));
        when(tasks.activeTasks()).thenReturn(Collections.singleton(task));
        taskManager.assignNextTask(taskExecutor);
        taskManager.setUncaughtException(exception, task.id());

        assertEquals(taskManager.drainUncaughtExceptions(), Collections.singletonMap(task.id(), exception));
    }

    @Test
    public void shouldClearExceptionsOnDrainExceptions() {
        taskManager.add(Collections.singleton(task));
        when(tasks.activeTasks()).thenReturn(Collections.singleton(task));
        taskManager.assignNextTask(taskExecutor);
        taskManager.setUncaughtException(exception, task.id());
        taskManager.drainUncaughtExceptions();

        assertEquals(taskManager.drainUncaughtExceptions(), Collections.emptyMap());
    }

    @Test
    public void shouldUnassignLockingTask() {
        final KafkaFutureImpl<StreamTask> future = new KafkaFutureImpl<>();

        taskManager.add(Collections.singleton(task));
        when(tasks.activeTasks()).thenReturn(Collections.singleton(task));
        when(tasks.task(task.id())).thenReturn(task);
        when(tasks.contains(task.id())).thenReturn(true);
        when(taskExecutor.unassign()).thenReturn(future);

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

        final KafkaFuture<Void> lockFuture = taskManager.lockAllTasks();
        assertFalse(lockFuture.isDone());

        verify(taskExecutor).unassign();

        future.complete(task);
        assertTrue(lockFuture.isDone());
    }
}


Review comment (suggestion):
I think, it would be OK to merge these two tests. Just add the `assertEquals()` of the second test to the end of the first test.

Code after:
        when(tasks.activeTasks()).thenReturn(Collections.singleton(task));
        taskManager.assignNextTask(taskExecutor);
        taskManager.setUncaughtException(exception, task.id());

        final Exception e = assertThrows(IllegalArgumentException.class, () -> taskManager.setUncaughtException(exception, task.id()));
        assertEquals("The uncaught exception must be cleared before restarting processing", e.getMessage());
    }

    @Test
    public void shouldReturnAndClearExceptionsOnDrainExceptions() {
        taskManager.add(Collections.singleton(task));
        when(tasks.activeTasks()).thenReturn(Collections.singleton(task));
        taskManager.assignNextTask(taskExecutor);
        taskManager.setUncaughtException(exception, task.id());

        assertEquals(taskManager.drainUncaughtExceptions(), Collections.singletonMap(task.id(), exception));
        assertEquals(taskManager.drainUncaughtExceptions(), Collections.emptyMap());
    }

    @Test
    public void shouldUnassignLockingTask() {
        final KafkaFutureImpl<StreamTask> future = new KafkaFutureImpl<>();

        taskManager.add(Collections.singleton(task));
        when(tasks.activeTasks()).thenReturn(Collections.singleton(task));
        when(tasks.task(task.id())).thenReturn(task);
        when(tasks.contains(task.id())).thenReturn(true);
        when(taskExecutor.unassign()).thenReturn(future);

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

        final KafkaFuture<Void> lockFuture = taskManager.lockAllTasks();
        assertFalse(lockFuture.isDone());

        verify(taskExecutor).unassign();

        future.complete(task);
        assertTrue(lockFuture.isDone());
    }
}