@Test
    public void testWakeupFromAssignmentCallback() {
        final String topic = "topic1";
        TopicPartition partition = new TopicPartition(topic, 0);
        final String consumerId = "follower";
        Set<String> topics = Collections.singleton(topic);
        MockRebalanceListener rebalanceListener = new MockRebalanceListener() {
            @Override
            public void onPartitionsAssigned(Collection<TopicPartition> partitions) {
                boolean raiseWakeup = this.assignedCount == 0;
                super.onPartitionsAssigned(partitions);

                if (raiseWakeup)
                    throw new WakeupException();
            }
        };

        subscriptions.subscribe(topics, Optional.of(rebalanceListener));

        // we only have metadata for one topic initially
        client.updateMetadata(RequestTestUtils.metadataUpdateWith(1, singletonMap(topic1, 1)));

        client.prepareResponse(groupCoordinatorResponse(node, Errors.NONE));
        coordinator.ensureCoordinatorReady(time.timer(Long.MAX_VALUE));

        // prepare initial rebalance
        partitionAssignor.prepare(singletonMap(consumerId, Collections.singletonList(partition)));

        client.prepareResponse(joinGroupFollowerResponse(1, consumerId, "leader", Errors.NONE));
        client.prepareResponse(syncGroupResponse(Collections.singletonList(partition), Errors.NONE));

        // The first call to poll should raise the exception from the rebalance listener
        try {
            coordinator.poll(time.timer(Long.MAX_VALUE));
            fail("Expected exception thrown from assignment callback");
        } catch (WakeupException e) {
        }

        // The second call should retry the assignment callback and succeed
        coordinator.poll(time.timer(Long.MAX_VALUE));

        assertFalse(coordinator.rejoinNeededOrPending());
        assertEquals(0, rebalanceListener.revokedCount);
        assertEquals(2, rebalanceListener.assignedCount);
    }
Verifies that subscription change updates SubscriptionState correctly even after JoinGroup failures
that don't re-invoke onJoinPrepare.
clients/src/test/java/org/apache/kafka/clients/consumer/internals/ConsumerCoordinatorTest.java