@Override
    protected Map<String, ByteBuffer> onLeaderElected(String leaderId,
                                                      String assignmentStrategy,
                                                      List<JoinGroupResponseData.JoinGroupResponseMember> allSubscriptions,
                                                      boolean skipAssignment) {
        ConsumerPartitionAssignor assignor = lookupAssignor(assignmentStrategy);
        if (assignor == null)
            throw new IllegalStateException("Coordinator selected invalid assignment protocol: " + assignmentStrategy);
        String assignorName = assignor.name();

        Set<String> allSubscribedTopics = new HashSet<>();
        Map<String, Subscription> subscriptions = new HashMap<>();

        // collect all the owned partitions
        Map<String, List<TopicPartition>> ownedPartitions = new HashMap<>();

        for (JoinGroupResponseData.JoinGroupResponseMember memberSubscription : allSubscriptions) {
            Subscription subscription = ConsumerProtocol.deserializeSubscription(ByteBuffer.wrap(memberSubscription.metadata()));
            subscription.setGroupInstanceId(Optional.ofNullable(memberSubscription.groupInstanceId()));
            subscriptions.put(memberSubscription.memberId(), subscription);
            allSubscribedTopics.addAll(subscription.topics());
            ownedPartitions.put(memberSubscription.memberId(), subscription.ownedPartitions());
        }

        // the leader will begin watching for changes to any of the topics the group is interested in,
        // which ensures that all metadata changes will eventually be seen
        updateGroupSubscription(allSubscribedTopics);

        isLeader = true;

        if (skipAssignment) {
            log.info("Skipped assignment for returning static leader at generation {}. The static leader " +
                "will continue with its existing assignment.", generation().generationId);
            assignmentSnapshot = metadataSnapshot;
            return Collections.emptyMap();
        }

        log.debug("Performing assignment using strategy {} with subscriptions {}", assignorName, subscriptions);

        Map<String, Assignment> assignments = assignor.assign(metadata.fetch(), new GroupSubscription(subscriptions)).groupAssignment();

        // skip the validation for built-in cooperative sticky assignor since we've considered
        // the "generation" of ownedPartition inside the assignor
        if (protocol == RebalanceProtocol.COOPERATIVE && !assignorName.equals(COOPERATIVE_STICKY_ASSIGNOR_NAME)) {
            validateCooperativeAssignment(ownedPartitions, assignments);
        }

        maybeUpdateGroupSubscription(assignorName, assignments, allSubscribedTopics);

        // metadataSnapshot could be updated when the subscription is updated therefore
        // we must take the assignment snapshot after.
        assignmentSnapshot = metadataSnapshot;

        log.info("Finished assignment for group at generation {}: {}", generation().generationId, assignments);

        Map<String, ByteBuffer> groupAssignment = new HashMap<>();
        for (Map.Entry<String, Assignment> assignmentEntry : assignments.entrySet()) {
            ByteBuffer buffer = ConsumerProtocol.serializeAssignment(assignmentEntry.getValue());
            groupAssignment.put(assignmentEntry.getKey(), buffer);
        }

        return groupAssignment;
    }
@SuppressWarnings("unchecked")
    @Test
    public void testPerformAssignmentShouldUpdateGroupSubscriptionAfterAssignmentIfNeeded() {
        SubscriptionState mockSubscriptionState = Mockito.mock(SubscriptionState.class);

        // the consumer only subscribed to "topic1"
        Map<String, List<String>> memberSubscriptions = singletonMap(consumerId, singletonList(topic1));

        List<JoinGroupResponseData.JoinGroupResponseMember> metadata = new ArrayList<>();
        for (Map.Entry<String, List<String>> subscriptionEntry : memberSubscriptions.entrySet()) {
            ConsumerPartitionAssignor.Subscription subscription = new ConsumerPartitionAssignor.Subscription(subscriptionEntry.getValue());
            ByteBuffer buf = ConsumerProtocol.serializeSubscription(subscription);
            metadata.add(new JoinGroupResponseData.JoinGroupResponseMember()
                .setMemberId(subscriptionEntry.getKey())
                .setMetadata(buf.array()));
        }

        // normal case: the assignment result will have partitions for only the subscribed topic: "topic1"
        partitionAssignor.prepare(Collections.singletonMap(consumerId, singletonList(t1p)));

        try (ConsumerCoordinator coordinator = buildCoordinator(rebalanceConfig, new Metrics(), assignors, false, mockSubscriptionState)) {
            coordinator.onLeaderElected("1", partitionAssignor.name(), metadata, false);

            ArgumentCaptor<Collection<String>> topicsCaptor = ArgumentCaptor.forClass(Collection.class);
            // groupSubscribe should be only called 1 time, which is before assignment,
            // because the assigned topics are the same as the subscribed topics
            Mockito.verify(mockSubscriptionState, Mockito.times(1)).groupSubscribe(topicsCaptor.capture());

            List<Collection<String>> capturedTopics = topicsCaptor.getAllValues();

            // expected the final group subscribed topics to be updated to "topic1"
            assertEquals(Collections.singleton(topic1), capturedTopics.get(0));
        }

        Mockito.clearInvocations(mockSubscriptionState);

        // unsubscribed topic partition assigned case: the assignment result will have partitions for (1) subscribed topic: "topic1"
        // and (2) the additional unsubscribed topic: "topic2". We should add "topic2" into group subscription list
        partitionAssignor.prepare(Collections.singletonMap(consumerId, Arrays.asList(t1p, t2p)));

        try (ConsumerCoordinator coordinator = buildCoordinator(rebalanceConfig, new Metrics(), assignors, false, mockSubscriptionState)) {
            coordinator.onLeaderElected("1", partitionAssignor.name(), metadata, false);

            ArgumentCaptor<Collection<String>> topicsCaptor = ArgumentCaptor.forClass(Collection.class);
            // groupSubscribe should be called 2 times, once before assignment, once after assignment
            // (because the assigned topics are not the same as the subscribed topics)
            Mockito.verify(mockSubscriptionState, Mockito.times(2)).groupSubscribe(topicsCaptor.capture());

            List<Collection<String>> capturedTopics = topicsCaptor.getAllValues();

            // expected the final group subscribed topics to be updated to "topic1" and "topic2"
            Set<String> expectedTopicsGotCalled = new HashSet<>(Arrays.asList(topic1, topic2));
            assertEquals(expectedTopicsGotCalled, capturedTopics.get(1));
        }
    }