void handleRebalanceComplete() {
        // we should pause consumer only within the listener since
        // before then the assignment has not been updated yet.
        if (stateUpdater == null) {
            mainConsumer.pause(mainConsumer.assignment());
        } else {
            // All tasks that are owned by the task manager are ready and do not need to be paused
            final Set<TopicPartition> partitionsNotToPause = tasks.allTasks()
                .stream()
                .flatMap(task -> task.inputPartitions().stream())
                .collect(Collectors.toSet());
            final Set<TopicPartition> partitionsToPause = new HashSet<>(mainConsumer.assignment());
            partitionsToPause.removeAll(partitionsNotToPause);
            mainConsumer.pause(partitionsToPause);
        }

        releaseLockedUnassignedTaskDirectories();

        rebalanceInProgress = false;
    }
@Test
    public void shouldPauseAllTopicsWithoutStateUpdaterOnRebalanceComplete() {
        final Set<TopicPartition> assigned = mkSet(t1p0, t1p1);
        when(consumer.assignment()).thenReturn(assigned);

        taskManager.handleRebalanceComplete();

        verify(consumer).pause(assigned);
    }