@Override
    public void onPartitionsRevoked(final Collection<TopicPartition> partitions) {
        log.debug("Current state {}: revoked partitions {} because of consumer rebalance.\n" +
                      "\tcurrently assigned active tasks: {}\n" +
                      "\tcurrently assigned standby tasks: {}\n",
                  streamThread.state(),
                  partitions,
                  taskManager.activeTaskIds(),
                  taskManager.standbyTaskIds());

        // We need to still invoke handleRevocation if the thread has been told to shut down, but we shouldn't ever
        // transition away from PENDING_SHUTDOWN once it's been initiated (to anything other than DEAD)
        if ((streamThread.setState(State.PARTITIONS_REVOKED) != null || streamThread.state() == State.PENDING_SHUTDOWN) && !partitions.isEmpty()) {
            final long start = time.milliseconds();
            try {
                taskManager.handleRevocation(partitions);
            } finally {
                log.info("partition revocation took {} ms.", time.milliseconds() - start);
            }
        }
    }
@Test
    public void shouldNotHandleRevokedPartitionsIfStateCannotTransitToPartitionRevoked() {
        when(streamThread.setState(State.PARTITIONS_REVOKED)).thenReturn(null);

        streamsRebalanceListener.onPartitionsRevoked(Collections.singletonList(new TopicPartition("topic", 0)));

        verify(taskManager, never()).handleRevocation(any());
    }