@Test
    public void testStaleMemberWaitsForCallbackToRejoinWhenTimerReset() {
        MembershipManagerImpl membershipManager = createMemberInStableState();
        Uuid topicId = Uuid.randomUuid();
        String topicName = "topic1";
        int ownedPartition = 0;
        TopicPartition tp = new TopicPartition(topicName, ownedPartition);
        mockOwnedPartitionAndAssignmentReceived(membershipManager, topicId, topicName,
            Collections.singletonList(new TopicIdPartition(topicId, tp)));
        CounterConsumerRebalanceListener listener = new CounterConsumerRebalanceListener();
        ConsumerRebalanceListenerInvoker invoker = consumerRebalanceListenerInvoker();
        when(subscriptionState.rebalanceListener()).thenReturn(Optional.of(listener));

        membershipManager.transitionToSendingLeaveGroup(true);
        membershipManager.onHeartbeatRequestGenerated();

        assertEquals(MemberState.STALE, membershipManager.state());
        assertFalse(backgroundEventQueue.isEmpty());
        assertInstanceOf(ConsumerRebalanceListenerCallbackNeededEvent.class, backgroundEventQueue.peek());

        // Stale member triggers onPartitionLost callback that will not complete just yet
        ConsumerRebalanceListenerCallbackCompletedEvent callbackEvent = performCallback(
            membershipManager,
            invoker,
            ConsumerRebalanceListenerMethodName.ON_PARTITIONS_LOST,
            topicPartitions(topicName, ownedPartition),
            false
        );

        // Timer reset while callback hasn't completed. Member should stay in STALE while it
        // completes releasing its assignment, and then transition to joining.
        membershipManager.maybeRejoinStaleMember();
        assertEquals(MemberState.STALE, membershipManager.state(), "Member should not transition " +
            "out of the STALE state when the timer is reset if the callback has not completed.");
        // Member should not clear its assignment to rejoin until the callback completes
        verify(subscriptionState, never()).assignFromSubscribed(any());

        completeCallback(callbackEvent, membershipManager);
        assertEquals(MemberState.JOINING, membershipManager.state());
        verify(subscriptionState).assignFromSubscribed(Collections.emptySet());
    }
This test should be the case when an assignment is sent to the member, and it cannot find
it in metadata (temporarily). If the broker continues to send the assignment to the
member, this one should keep it waiting for metadata and continue to request updates.
clients/src/test/java/org/apache/kafka/clients/consumer/internals/MembershipManagerImplTest.java