@Override
    public void onPartitionsAssigned(final Collection<TopicPartition> partitions) {
        // NB: all task management is already handled by:
        // org.apache.kafka.streams.processor.internals.StreamsPartitionAssignor.onAssignment
        if (assignmentErrorCode.get() == AssignorError.INCOMPLETE_SOURCE_TOPIC_METADATA.code()) {
            log.error("Received error code {}. {}",
                    AssignorError.INCOMPLETE_SOURCE_TOPIC_METADATA.codeName(),
                    AssignorError.INCOMPLETE_SOURCE_TOPIC_METADATA.description());
            taskManager.handleRebalanceComplete();
            throw new MissingSourceTopicException("One or more source topics were missing during rebalance");
        } else if (assignmentErrorCode.get() == AssignorError.VERSION_PROBING.code()) {
            log.info("Received version probing code {}", AssignorError.VERSION_PROBING);
        }  else if (assignmentErrorCode.get() == AssignorError.ASSIGNMENT_ERROR.code()) {
            log.error("Received error code {}", AssignorError.ASSIGNMENT_ERROR);
            taskManager.handleRebalanceComplete();
            throw new TaskAssignmentException("Hit an unexpected exception during task assignment phase of rebalance");
        } else if (assignmentErrorCode.get() == AssignorError.SHUTDOWN_REQUESTED.code()) {
            log.error("A Kafka Streams client in this Kafka Streams application is requesting to shutdown the application");
            taskManager.handleRebalanceComplete();
            streamThread.shutdownToError();
            return;
        } else if (assignmentErrorCode.get() != AssignorError.NONE.code()) {
            log.error("Received unknown error code {}", assignmentErrorCode.get());
            throw new TaskAssignmentException("Hit an unrecognized exception during rebalance");
        }

        streamThread.setState(State.PARTITIONS_ASSIGNED);
        streamThread.setPartitionAssignedTime(time.milliseconds());
        taskManager.handleRebalanceComplete();
    }
@Test
    public void shouldHandleAssignedPartitions() {
        assignmentErrorCode.set(AssignorError.NONE.code());

        streamsRebalanceListener.onPartitionsAssigned(Collections.emptyList());

        verify(streamThread).setState(State.PARTITIONS_ASSIGNED);
        verify(streamThread).setPartitionAssignedTime(time.milliseconds());
        verify(taskManager).handleRebalanceComplete();
    }