@Test
public void testWakeupDuringJoin() {

    final String consumerId = "leader";
    subscriptions.subscribe(Arrays.asList(topicName), rebalanceListener);
    subscriptions.needReassignment();
    // ensure metadata is up-to-date for leader
    metadata.setTopics(Arrays.asList(topicName));
    metadata.update(cluster, time.milliseconds());
    client.prepareResponse(consumerMetadataResponse(node, Errors.NONE.code()));
    coordinator.ensureCoordinatorReady();
    Map<String, List<String>> memberSubscriptions = Collections.singletonMap(consumerId, Arrays.asList(topicName));
    partitionAssignor.prepare(Collections.singletonMap(consumerId, Arrays.asList(tp)));
    // prepare only the first half of the join and then trigger the wakeup
    client.prepareResponse(joinGroupLeaderResponse(1, consumerId, memberSubscriptions, Errors.NONE.code()));
    consumerClient.wakeup();
    try {
        coordinator.ensurePartitionAssignment();
    } catch (WakeupException e) {
        // ignore
    }
    // now complete the second half
    client.prepareResponse(syncGroupResponse(Arrays.asList(tp), Errors.NONE.code()));
    coordinator.ensurePartitionAssignment();
    assertFalse(subscriptions.partitionAssignmentNeeded());
    assertEquals(Collections.singleton(tp), subscriptions.assignedPartitions());
    assertEquals(1, rebalanceListener.revokedCount);
    assertEquals(Collections.emptySet(), rebalanceListener.revoked);
    assertEquals(1, rebalanceListener.assignedCount);
    assertEquals(Collections.singleton(tp), rebalanceListener.assigned);
}