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

Code before:
        assertTrue(awaitingRunnable.awaitDone.await(VERIFICATION_TIMEOUT, TimeUnit.MILLISECONDS));

        awaitingRunnable.shutdown();
    }

    @Test
    public void shouldReturnFromAwaitOnSignalProcessableTasks() throws InterruptedException {
        final AwaitingRunnable awaitingRunnable = new AwaitingRunnable();
        final Thread awaitingThread = new Thread(awaitingRunnable);
        awaitingThread.start();
        verify(tasks, timeout(VERIFICATION_TIMEOUT).atLeastOnce()).activeTasks();

        taskManager.signalProcessableTasks();

        assertTrue(awaitingRunnable.awaitDone.await(VERIFICATION_TIMEOUT, TimeUnit.MILLISECONDS));

        awaitingRunnable.shutdown();
    }

    @Test
    public void shouldReturnFromAwaitOnUnassignment() throws InterruptedException {
        taskManager.add(Collections.singleton(task));
        when(taskExecutionMetadata.canProcessTask(eq(task), anyLong())).thenReturn(true);

        final StreamTask t = taskManager.assignNextTask(taskExecutor);
        assertNotNull(t);
        final AwaitingRunnable awaitingRunnable = new AwaitingRunnable();
        final Thread awaitingThread = new Thread(awaitingRunnable);
        awaitingThread.start();
        verify(tasks, timeout(VERIFICATION_TIMEOUT).atLeastOnce()).activeTasks();

        taskManager.unassignTask(t, taskExecutor);

        assertTrue(awaitingRunnable.awaitDone.await(VERIFICATION_TIMEOUT, TimeUnit.MILLISECONDS));

        awaitingRunnable.shutdown();
    }

    @Test
    public void shouldReturnFromAwaitOnAdding() throws InterruptedException {
        final AwaitingRunnable awaitingRunnable = new AwaitingRunnable();
        final Thread awaitingThread = new Thread(awaitingRunnable);
        awaitingThread.start();
        verify(tasks, timeout(VERIFICATION_TIMEOUT).atLeastOnce()).activeTasks();

        taskManager.add(Collections.singleton(task));

        assertTrue(awaitingRunnable.awaitDone.await(VERIFICATION_TIMEOUT, TimeUnit.MILLISECONDS));

        awaitingRunnable.shutdown();
    }

Review comment (suggestion):
nit:
```suggestion
        final StreamTask task = taskManager.assignNextTask(taskExecutor);
        assertNotNull(task);
```

Code after:
        assertTrue(awaitingRunnable.awaitDone.await(VERIFICATION_TIMEOUT, TimeUnit.MILLISECONDS));

        awaitingRunnable.shutdown();
    }

    @Test
    public void shouldReturnFromAwaitOnSignalProcessableTasks() throws InterruptedException {
        final AwaitingRunnable awaitingRunnable = new AwaitingRunnable();
        final Thread awaitingThread = new Thread(awaitingRunnable);
        awaitingThread.start();
        verify(tasks, timeout(VERIFICATION_TIMEOUT).atLeastOnce()).activeTasks();

        taskManager.signalProcessableTasks();

        assertTrue(awaitingRunnable.awaitDone.await(VERIFICATION_TIMEOUT, TimeUnit.MILLISECONDS));

        awaitingRunnable.shutdown();
    }

    @Test
    public void shouldReturnFromAwaitOnUnassignment() throws InterruptedException {
        taskManager.add(Collections.singleton(task));
        when(taskExecutionMetadata.canProcessTask(eq(task), anyLong())).thenReturn(true);

        final StreamTask task = taskManager.assignNextTask(taskExecutor);
        assertNotNull(task);
        final AwaitingRunnable awaitingRunnable = new AwaitingRunnable();
        final Thread awaitingThread = new Thread(awaitingRunnable);
        awaitingThread.start();
        verify(tasks, timeout(VERIFICATION_TIMEOUT).atLeastOnce()).activeTasks();

        taskManager.unassignTask(task, taskExecutor);

        assertTrue(awaitingRunnable.awaitDone.await(VERIFICATION_TIMEOUT, TimeUnit.MILLISECONDS));

        awaitingRunnable.shutdown();
    }

    @Test
    public void shouldReturnFromAwaitOnAdding() throws InterruptedException {
        final AwaitingRunnable awaitingRunnable = new AwaitingRunnable();
        final Thread awaitingThread = new Thread(awaitingRunnable);
        awaitingThread.start();
        verify(tasks, timeout(VERIFICATION_TIMEOUT).atLeastOnce()).activeTasks();

        taskManager.add(Collections.singleton(task));

        assertTrue(awaitingRunnable.awaitDone.await(VERIFICATION_TIMEOUT, TimeUnit.MILLISECONDS));

        awaitingRunnable.shutdown();
    }