/*
 * Licensed to the Apache Software Foundation (ASF) under one or more
 * contributor license agreements. See the NOTICE file distributed with
 * this work for additional information regarding copyright ownership.
 * The ASF licenses this file to You under the Apache License, Version 2.0
 * (the "License"); you may not use this file except in compliance with
 * the License. You may obtain a copy of the License at
 *
 *    http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package org.apache.kafka.clients.consumer;

import org.apache.kafka.clients.ClientRequest;
import org.apache.kafka.clients.CommonClientConfigs;
import org.apache.kafka.clients.KafkaClient;
import org.apache.kafka.clients.MockClient;
import org.apache.kafka.clients.NetworkClient;
import org.apache.kafka.clients.NodeApiVersions;
import org.apache.kafka.clients.consumer.internals.AsyncKafkaConsumer;
import org.apache.kafka.clients.consumer.internals.AutoOffsetResetStrategy;
import org.apache.kafka.clients.consumer.internals.ClassicKafkaConsumer;
import org.apache.kafka.clients.consumer.internals.ConsumerMetadata;
import org.apache.kafka.clients.consumer.internals.ConsumerProtocol;
import org.apache.kafka.clients.consumer.internals.MockRebalanceListener;
import org.apache.kafka.clients.consumer.internals.SubscriptionState;
import org.apache.kafka.common.Cluster;
import org.apache.kafka.common.IsolationLevel;
import org.apache.kafka.common.KafkaException;
import org.apache.kafka.common.Metric;
import org.apache.kafka.common.MetricName;
import org.apache.kafka.common.Node;
import org.apache.kafka.common.TopicIdPartition;
import org.apache.kafka.common.TopicPartition;
import org.apache.kafka.common.Uuid;
import org.apache.kafka.common.compress.Compression;
import org.apache.kafka.common.config.SaslConfigs;
import org.apache.kafka.common.config.SslConfigs;
import org.apache.kafka.common.errors.AuthenticationException;
import org.apache.kafka.common.errors.InterruptException;
import org.apache.kafka.common.errors.InvalidConfigurationException;
import org.apache.kafka.common.errors.InvalidGroupIdException;
import org.apache.kafka.common.errors.InvalidTopicException;
import org.apache.kafka.common.errors.RecordDeserializationException;
import org.apache.kafka.common.errors.SerializationException;
import org.apache.kafka.common.errors.UnsupportedVersionException;
import org.apache.kafka.common.errors.WakeupException;
import org.apache.kafka.common.header.Headers;
import org.apache.kafka.common.internals.ClusterResourceListeners;
import org.apache.kafka.common.message.FetchResponseData;
import org.apache.kafka.common.message.HeartbeatResponseData;
import org.apache.kafka.common.message.JoinGroupRequestData;
import org.apache.kafka.common.message.JoinGroupResponseData;
import org.apache.kafka.common.message.LeaveGroupResponseData;
import org.apache.kafka.common.message.ListOffsetsRequestData.ListOffsetsPartition;
import org.apache.kafka.common.message.ListOffsetsResponseData;
import org.apache.kafka.common.message.ListOffsetsResponseData.ListOffsetsPartitionResponse;
import org.apache.kafka.common.message.ListOffsetsResponseData.ListOffsetsTopicResponse;
import org.apache.kafka.common.message.OffsetFetchResponseData;
import org.apache.kafka.common.message.SyncGroupResponseData;
import org.apache.kafka.common.metrics.JmxReporter;
import org.apache.kafka.common.metrics.KafkaMetric;
import org.apache.kafka.common.metrics.Measurable;
import org.apache.kafka.common.metrics.MetricConfig;
import org.apache.kafka.common.metrics.Metrics;
import org.apache.kafka.common.metrics.Monitorable;
import org.apache.kafka.common.metrics.PluginMetrics;
import org.apache.kafka.common.metrics.Sensor;
import org.apache.kafka.common.metrics.stats.Avg;
import org.apache.kafka.common.network.Selectable;
import org.apache.kafka.common.protocol.ApiKeys;
import org.apache.kafka.common.protocol.Errors;
import org.apache.kafka.common.record.MemoryRecords;
import org.apache.kafka.common.record.MemoryRecordsBuilder;
import org.apache.kafka.common.record.TimestampType;
import org.apache.kafka.common.requests.AbstractResponse;
import org.apache.kafka.common.requests.FetchRequest;
import org.apache.kafka.common.requests.FetchResponse;
import org.apache.kafka.common.requests.FindCoordinatorResponse;
import org.apache.kafka.common.requests.HeartbeatResponse;
import org.apache.kafka.common.requests.JoinGroupRequest;
import org.apache.kafka.common.requests.JoinGroupResponse;
import org.apache.kafka.common.requests.LeaveGroupResponse;
import org.apache.kafka.common.requests.ListOffsetsRequest;
import org.apache.kafka.common.requests.ListOffsetsResponse;
import org.apache.kafka.common.requests.MetadataResponse;
import org.apache.kafka.common.requests.OffsetCommitRequest;
import org.apache.kafka.common.requests.OffsetCommitResponse;
import org.apache.kafka.common.requests.OffsetFetchResponse;
import org.apache.kafka.common.requests.RequestTestUtils;
import org.apache.kafka.common.requests.SyncGroupResponse;
import org.apache.kafka.common.security.auth.SecurityProtocol;
import org.apache.kafka.common.serialization.ByteArrayDeserializer;
import org.apache.kafka.common.serialization.Deserializer;
import org.apache.kafka.common.serialization.StringDeserializer;
import org.apache.kafka.common.telemetry.internals.ClientTelemetryReporter;
import org.apache.kafka.common.telemetry.internals.ClientTelemetrySender;
import org.apache.kafka.common.utils.LogCaptureAppender;
import org.apache.kafka.common.utils.LogContext;
import org.apache.kafka.common.utils.MockTime;
import org.apache.kafka.common.utils.Time;
import org.apache.kafka.common.utils.Utils;
import org.apache.kafka.test.MockConsumerInterceptor;
import org.apache.kafka.test.MockDeserializer;
import org.apache.kafka.test.MockMetricsReporter;
import org.apache.kafka.test.TestUtils;

import org.apache.logging.log4j.Level;
import org.junit.jupiter.api.AfterEach;
import org.junit.jupiter.api.Timeout;
import org.junit.jupiter.params.ParameterizedTest;
import org.junit.jupiter.params.provider.EnumSource;
import org.mockito.MockedStatic;
import org.mockito.Mockito;
import org.mockito.internal.stubbing.answers.CallsRealMethods;

import java.lang.management.ManagementFactory;
import java.nio.ByteBuffer;
import java.time.Duration;
import java.util.AbstractMap;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collection;
import java.util.Collections;
import java.util.ConcurrentModificationException;
import java.util.HashMap;
import java.util.HashSet;
import java.util.Iterator;
import java.util.LinkedHashMap;
import java.util.List;
import java.util.Locale;
import java.util.Map;
import java.util.Optional;
import java.util.OptionalLong;
import java.util.Properties;
import java.util.Queue;
import java.util.Set;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.Future;
import java.util.concurrent.ScheduledExecutorService;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.TimeoutException;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.concurrent.atomic.AtomicReference;
import java.util.regex.Pattern;
import java.util.stream.Collectors;
import java.util.stream.Stream;

import javax.management.MBeanServer;
import javax.management.ObjectName;
import javax.security.auth.login.LoginException;

import static java.util.Collections.singletonList;
import static org.apache.kafka.clients.consumer.internals.ClassicKafkaConsumer.DEFAULT_REASON;
import static org.apache.kafka.clients.consumer.internals.ConsumerUtils.CONSUMER_METRIC_GROUP_PREFIX;
import static org.apache.kafka.clients.consumer.internals.ConsumerUtils.COORDINATOR_METRICS_SUFFIX;
import static org.apache.kafka.common.requests.FetchMetadata.INVALID_SESSION_ID;
import static org.apache.kafka.common.utils.Utils.propsToMap;
import static org.junit.jupiter.api.Assertions.assertDoesNotThrow;
import static org.junit.jupiter.api.Assertions.assertEquals;
import static org.junit.jupiter.api.Assertions.assertFalse;
import static org.junit.jupiter.api.Assertions.assertInstanceOf;
import static org.junit.jupiter.api.Assertions.assertNotEquals;
import static org.junit.jupiter.api.Assertions.assertNotNull;
import static org.junit.jupiter.api.Assertions.assertNotSame;
import static org.junit.jupiter.api.Assertions.assertNull;
import static org.junit.jupiter.api.Assertions.assertThrows;
import static org.junit.jupiter.api.Assertions.assertTrue;
import static org.junit.jupiter.api.Assertions.fail;
import static org.mockito.ArgumentMatchers.any;
import static org.mockito.ArgumentMatchers.anyString;
import static org.mockito.Mockito.atMostOnce;
import static org.mockito.Mockito.mock;
import static org.mockito.Mockito.mockStatic;
import static org.mockito.Mockito.never;
import static org.mockito.Mockito.spy;
import static org.mockito.Mockito.verify;
import static org.mockito.Mockito.when;

/**
 * Note to future authors in this class. If you close the consumer, close with DURATION.ZERO to reduce the duration of
 * the test.
 */
@Timeout(value = 120)
public class KafkaConsumerTest {

    private final String topic = "test";
    private final Uuid topicId = Uuid.randomUuid();
    private final TopicPartition tp0 = new TopicPartition(topic, 0);
    private final TopicPartition tp1 = new TopicPartition(topic, 1);

    private final String topic2 = "test2";
    private final Uuid topicId2 = Uuid.randomUuid();
    private final TopicPartition t2p0 = new TopicPartition(topic2, 0);

    private final String topic3 = "test3";
    private final Uuid topicId3 = Uuid.randomUuid();
    private final TopicPartition t3p0 = new TopicPartition(topic3, 0);

    private final int sessionTimeoutMs = 10000;
    private final int defaultApiTimeoutMs = 60000;
    private final int heartbeatIntervalMs = 1000;

    // Set auto commit interval lower than heartbeat so we don't need to deal with
    // a concurrent heartbeat request
    private final int autoCommitIntervalMs = 500;

    private final String groupId = "mock-group";
    private final String memberId = "memberId";
    private final String leaderId = "leaderId";
    private final Optional<String> groupInstanceId = Optional.of("mock-instance");
    private final Map<String, Uuid> topicIds = Stream.of(
            new AbstractMap.SimpleEntry<>(topic, topicId),
            new AbstractMap.SimpleEntry<>(topic2, topicId2),
            new AbstractMap.SimpleEntry<>(topic3, topicId3))
            .collect(Collectors.toMap(Map.Entry::getKey, Map.Entry::getValue));
    private final Map<Uuid, String> topicNames = Stream.of(
            new AbstractMap.SimpleEntry<>(topicId, topic),
            new AbstractMap.SimpleEntry<>(topicId2, topic2),
            new AbstractMap.SimpleEntry<>(topicId3, topic3))
            .collect(Collectors.toMap(Map.Entry::getKey, Map.Entry::getValue));

    private final String partitionRevoked = "Hit partition revoke ";
    private final String partitionAssigned = "Hit partition assign ";
    private final String partitionLost = "Hit partition lost ";

    private final Collection<TopicPartition> singleTopicPartition = Set.of(new TopicPartition(topic, 0));
    private final Time time = new MockTime();
    private final SubscriptionState subscription = new SubscriptionState(new LogContext(), AutoOffsetResetStrategy.EARLIEST);
    private final ConsumerPartitionAssignor assignor = new RoundRobinAssignor();

    private KafkaConsumer<?, ?> consumer;

    @AfterEach
    public void cleanup() {
        if (consumer != null) {
            consumer.close(CloseOptions.timeout(Duration.ZERO));
        }
    }

    @ParameterizedTest
    @EnumSource(GroupProtocol.class)
    public void testSubscribingCustomMetricsDoesntAffectConsumerMetrics(GroupProtocol groupProtocol) {
        Properties props = new Properties();
        props.setProperty(ConsumerConfig.GROUP_PROTOCOL_CONFIG, groupProtocol.name());
        props.setProperty(ConsumerConfig.BOOTSTRAP_SERVERS_CONFIG, "localhost:9999");
        consumer = newConsumer(props, new StringDeserializer(), new StringDeserializer());

        Map<MetricName, KafkaMetric> customMetrics = customMetrics();
        customMetrics.forEach((name, metric) -> consumer.registerMetricForSubscription(metric));

        Map<MetricName, ? extends Metric> consumerMetrics = consumer.metrics();
        customMetrics.forEach((name, metric) -> assertFalse(consumerMetrics.containsKey(name)));
    }

    @ParameterizedTest
    @EnumSource(GroupProtocol.class)
    public void testSubscribingCustomMetricsWithSameNameDoesntAffectConsumerMetrics(GroupProtocol groupProtocol) {
        Properties props = new Properties();
        props.setProperty(ConsumerConfig.GROUP_PROTOCOL_CONFIG, groupProtocol.name());
        props.setProperty(ConsumerConfig.BOOTSTRAP_SERVERS_CONFIG, "localhost:9999");
        Class<?> consumerClass = groupProtocol == GroupProtocol.CLASSIC ? ClassicKafkaConsumer.class : AsyncKafkaConsumer.class;
        try (final LogCaptureAppender appender = LogCaptureAppender.createAndRegister()) {
            appender.setClassLogger(consumerClass, Level.DEBUG);
            consumer = newConsumer(props, new StringDeserializer(), new StringDeserializer());
            KafkaMetric existingMetricToAdd = (KafkaMetric) consumer.metrics().entrySet().iterator().next().getValue();
            consumer.registerMetricForSubscription(existingMetricToAdd);
            final String expectedMessage = String.format("Skipping registration for metric %s. Existing consumer metrics cannot be overwritten.", existingMetricToAdd.metricName());
            assertTrue(appender.getMessages().stream().anyMatch(m -> m.contains(expectedMessage)));
        }
    }

    @ParameterizedTest
    @EnumSource(GroupProtocol.class)
    public void testAssignedPartitionsMetrics(GroupProtocol groupProtocol) throws InterruptedException {
        consumer = newConsumer(groupProtocol, time, mock(KafkaClient.class), subscription,
                mock(ConsumerMetadata.class), assignor, false, groupInstanceId);
        Metrics metrics = consumer.metricsRegistry();

        // This metric is added in the background thread for the AsyncConsumer, so waiting on it to avoid flakiness.
        TestUtils.waitForCondition(() -> getMetric(metrics, "assigned-partitions") != null,
                "Consumer should register the assigned-partitions metric");
        assertNotNull(getMetric(metrics, "assigned-partitions"));
        assertEquals(0.0d, getMetric(metrics, "assigned-partitions").metricValue());

        subscription.assignFromUser(Set.of(tp0));
        assertEquals(1.0d, getMetric(metrics, "assigned-partitions").metricValue());

        subscription.assignFromUser(Set.of(tp0, tp1));
        assertEquals(2.0d, getMetric(metrics, "assigned-partitions").metricValue());

        subscription.unsubscribe();
        subscription.subscribe(Set.of(topic), Optional.empty());
        subscription.assignFromSubscribed(Set.of(tp0));
        assertEquals(1.0d, getMetric(metrics, "assigned-partitions").metricValue());
    }

    private KafkaMetric getMetric(Metrics metrics, String name) {
        return metrics.metrics().get(metrics.metricName(name, CONSUMER_METRIC_GROUP_PREFIX + COORDINATOR_METRICS_SUFFIX));
    }

    @ParameterizedTest
    @EnumSource(GroupProtocol.class)
    public void testUnsubscribingCustomMetricsWithSameNameDoesntAffectConsumerMetrics(GroupProtocol groupProtocol) {
        Properties props = new Properties();
        props.setProperty(ConsumerConfig.GROUP_PROTOCOL_CONFIG, groupProtocol.name());
        props.setProperty(ConsumerConfig.BOOTSTRAP_SERVERS_CONFIG, "localhost:9999");
        Class<?> consumerClass = groupProtocol == GroupProtocol.CLASSIC ? ClassicKafkaConsumer.class : AsyncKafkaConsumer.class;
        try (final LogCaptureAppender appender = LogCaptureAppender.createAndRegister()) {
            appender.setClassLogger(consumerClass, Level.DEBUG);
            consumer = newConsumer(props, new StringDeserializer(), new StringDeserializer());
            KafkaMetric existingMetricToRemove = (KafkaMetric) consumer.metrics().entrySet().iterator().next().getValue();
            consumer.unregisterMetricFromSubscription(existingMetricToRemove);
            final String expectedMessage = String.format("Skipping unregistration for metric %s. Existing consumer metrics cannot be removed.", existingMetricToRemove.metricName());
            assertTrue(appender.getMessages().stream().anyMatch(m -> m.contains(expectedMessage)));
        }
    }

    @ParameterizedTest
    @EnumSource(GroupProtocol.class)
    public void testShouldOnlyCallMetricReporterMetricChangeOnceWithExistingConsumerMetric(GroupProtocol groupProtocol) {
        try (MockedStatic<CommonClientConfigs> mockedCommonClientConfigs = mockStatic(CommonClientConfigs.class, new CallsRealMethods())) {
            ClientTelemetryReporter clientTelemetryReporter = mock(ClientTelemetryReporter.class);
            clientTelemetryReporter.configure(any());
            mockedCommonClientConfigs.when(() -> CommonClientConfigs.telemetryReporter(anyString(), any())).thenReturn(Optional.of(clientTelemetryReporter));

            Properties props = new Properties();
            props.setProperty(ConsumerConfig.GROUP_PROTOCOL_CONFIG, groupProtocol.name());
            props.setProperty(ConsumerConfig.BOOTSTRAP_SERVERS_CONFIG, "localhost:9999");
            consumer = newConsumer(props, new StringDeserializer(), new StringDeserializer());

            KafkaMetric existingMetric = (KafkaMetric) consumer.metrics().entrySet().iterator().next().getValue();
            consumer.registerMetricForSubscription(existingMetric);
            // This test would fail without the check as the existing metric is registered in the consumer on startup
            Mockito.verify(clientTelemetryReporter, atMostOnce()).metricChange(existingMetric);
        }
    }

    @ParameterizedTest
    @EnumSource(GroupProtocol.class)
    public void testShouldNotCallMetricReporterMetricRemovalWithExistingConsumerMetric(GroupProtocol groupProtocol) {
        try (MockedStatic<CommonClientConfigs> mockedCommonClientConfigs = mockStatic(CommonClientConfigs.class, new CallsRealMethods())) {
            ClientTelemetryReporter clientTelemetryReporter = mock(ClientTelemetryReporter.class);
            clientTelemetryReporter.configure(any());
            mockedCommonClientConfigs.when(() -> CommonClientConfigs.telemetryReporter(anyString(), any())).thenReturn(Optional.of(clientTelemetryReporter));

            Properties props = new Properties();
            props.setProperty(ConsumerConfig.GROUP_PROTOCOL_CONFIG, groupProtocol.name());
            props.setProperty(ConsumerConfig.BOOTSTRAP_SERVERS_CONFIG, "localhost:9999");
            consumer = newConsumer(props, new StringDeserializer(), new StringDeserializer());

            KafkaMetric existingMetric = (KafkaMetric) consumer.metrics().entrySet().iterator().next().getValue();
            consumer.unregisterMetricFromSubscription(existingMetric);
            Mockito.verify(clientTelemetryReporter, never()).metricRemoval(existingMetric);
        }
    }

    @ParameterizedTest
    @EnumSource(GroupProtocol.class)
    public void testUnSubscribingNonExisingMetricsDoesntCauseError(GroupProtocol groupProtocol) {
        Properties props = new Properties();
        props.setProperty(ConsumerConfig.GROUP_PROTOCOL_CONFIG, groupProtocol.name());
        props.setProperty(ConsumerConfig.BOOTSTRAP_SERVERS_CONFIG, "localhost:9999");
        consumer = newConsumer(props, new StringDeserializer(), new StringDeserializer());

        Map<MetricName, KafkaMetric> customMetrics = customMetrics();
        //Metrics never registered but removed should not cause an error
        customMetrics.forEach((name, metric) -> assertDoesNotThrow(() -> consumer.unregisterMetricFromSubscription(metric)));
    }

    @ParameterizedTest
    @EnumSource(GroupProtocol.class)
    public void testMetricsReporterAutoGeneratedClientId(GroupProtocol groupProtocol) {
        Properties props = new Properties();
        props.setProperty(ConsumerConfig.GROUP_PROTOCOL_CONFIG, groupProtocol.name());
        props.setProperty(ConsumerConfig.BOOTSTRAP_SERVERS_CONFIG, "localhost:9999");
        props.setProperty(ConsumerConfig.METRIC_REPORTER_CLASSES_CONFIG, MockMetricsReporter.class.getName());
        consumer = newConsumer(props, new StringDeserializer(), new StringDeserializer());

        assertEquals(2, consumer.metricsRegistry().reporters().size());

        MockMetricsReporter mockMetricsReporter = (MockMetricsReporter) consumer.metricsRegistry().reporters().stream()
            .filter(reporter -> reporter instanceof MockMetricsReporter).findFirst().orElseThrow();
        assertEquals(consumer.clientId(), mockMetricsReporter.clientId);

        consumer.close(CloseOptions.timeout(Duration.ZERO));
    }

    @ParameterizedTest
    @EnumSource(GroupProtocol.class)
    public void testDisableJmxAndClientTelemetryReporter(GroupProtocol groupProtocol) {
        Properties props = new Properties();
        props.setProperty(ConsumerConfig.GROUP_PROTOCOL_CONFIG, groupProtocol.name());
        props.setProperty(ConsumerConfig.BOOTSTRAP_SERVERS_CONFIG, "localhost:9999");
        props.setProperty(ConsumerConfig.METRIC_REPORTER_CLASSES_CONFIG, "");
        props.setProperty(ConsumerConfig.ENABLE_METRICS_PUSH_CONFIG, "false");
        consumer = newConsumer(props, new StringDeserializer(), new StringDeserializer());
        assertTrue(consumer.metricsRegistry().reporters().isEmpty());
    }

    @ParameterizedTest
    @EnumSource(GroupProtocol.class)
    public void testExplicitlyOnlyEnableJmxReporter(GroupProtocol groupProtocol) {
        Properties props = new Properties();
        props.setProperty(ConsumerConfig.GROUP_PROTOCOL_CONFIG, groupProtocol.name());
        props.setProperty(ConsumerConfig.BOOTSTRAP_SERVERS_CONFIG, "localhost:9999");
        props.setProperty(ConsumerConfig.METRIC_REPORTER_CLASSES_CONFIG, "org.apache.kafka.common.metrics.JmxReporter");
        props.setProperty(ConsumerConfig.ENABLE_METRICS_PUSH_CONFIG, "false");
        consumer = newConsumer(props, new StringDeserializer(), new StringDeserializer());
        assertEquals(1, consumer.metricsRegistry().reporters().size());
        assertInstanceOf(JmxReporter.class, consumer.metricsRegistry().reporters().get(0));
    }

    @ParameterizedTest
    @EnumSource(GroupProtocol.class)
    public void testExplicitlyOnlyEnableClientTelemetryReporter(GroupProtocol groupProtocol) {
        Properties props = new Properties();
        props.setProperty(ConsumerConfig.GROUP_PROTOCOL_CONFIG, groupProtocol.name());
        props.setProperty(ConsumerConfig.BOOTSTRAP_SERVERS_CONFIG, "localhost:9999");
        props.setProperty(ConsumerConfig.METRIC_REPORTER_CLASSES_CONFIG, "");
        consumer = newConsumer(props, new StringDeserializer(), new StringDeserializer());
        assertEquals(1, consumer.metricsRegistry().reporters().size());
        assertInstanceOf(ClientTelemetryReporter.class, consumer.metricsRegistry().reporters().get(0));
    }

    // TODO: this test requires rebalance logic which is not yet implemented in the CONSUMER group protocol.
    //       Once it is implemented, this should use both group protocols.
    @ParameterizedTest
    @EnumSource(value = GroupProtocol.class, names = "CLASSIC")
    @SuppressWarnings("unchecked")
    public void testPollReturnsRecords(GroupProtocol groupProtocol) {
        consumer = setUpConsumerWithRecordsToPoll(groupProtocol, tp0, 5);

        ConsumerRecords<String, String> records = (ConsumerRecords<String, String>) consumer.poll(Duration.ZERO);

        assertEquals(5, records.count());
        assertEquals(Set.of(tp0), records.partitions());
        assertEquals(5, records.records(tp0).size());
        assertEquals(1, records.nextOffsets().size());
        assertEquals(new OffsetAndMetadata(5), records.nextOffsets().get(tp0));
    }

    // TODO: this test requires rebalance logic which is not yet implemented in the CONSUMER group protocol.
    //       Once it is implemented, this should use both group protocols.
    @ParameterizedTest
    @EnumSource(value = GroupProtocol.class, names = "CLASSIC")
    @SuppressWarnings("unchecked")
    public void testSecondPollWithDeserializationErrorThrowsRecordDeserializationException(GroupProtocol groupProtocol) {
        int invalidRecordNumber = 4;
        int invalidRecordOffset = 3;
        StringDeserializer deserializer = mockErrorDeserializer(invalidRecordNumber);

        consumer = setUpConsumerWithRecordsToPoll(groupProtocol, tp0, 5, deserializer);
        ConsumerRecords<String, String> records = (ConsumerRecords<String, String>) consumer.poll(Duration.ZERO);

        assertEquals(invalidRecordNumber - 1, records.count());
        assertEquals(Set.of(tp0), records.partitions());
        assertEquals(invalidRecordNumber - 1, records.records(tp0).size());
        long lastOffset = records.records(tp0).get(records.records(tp0).size() - 1).offset();
        assertEquals(invalidRecordNumber - 2, lastOffset);
        assertEquals(1, records.nextOffsets().size());
        assertEquals(new OffsetAndMetadata(lastOffset + 1), records.nextOffsets().get(tp0));

        RecordDeserializationException rde = assertThrows(RecordDeserializationException.class, () -> consumer.poll(Duration.ZERO));
        assertEquals(invalidRecordOffset, rde.offset());
        assertEquals(tp0, rde.topicPartition());
        assertEquals(consumer.position(tp0), rde.offset());
    }

    /*
        Create a mock deserializer which throws a SerializationException on the Nth record's value deserialization
     */
    private StringDeserializer mockErrorDeserializer(int recordNumber) {
        int recordIndex = recordNumber - 1;
        return new StringDeserializer() {
            int i = 0;
            @Override
            public String deserialize(String topic, byte[] data) {
                if (i == recordIndex) {
                    throw new SerializationException();
                } else {
                    i++;
                    return super.deserialize(topic, data);
                }
            }

            @Override
            public String deserialize(String topic, Headers headers, ByteBuffer data) {
                if (i == recordIndex) {
                    throw new SerializationException();
                } else {
                    i++;
                    return super.deserialize(topic, headers, data);
                }
            }
        };
    }

    private KafkaConsumer<?, ?> setUpConsumerWithRecordsToPoll(GroupProtocol groupProtocol,
                                                               TopicPartition tp,
                                                               int recordCount) {
        return setUpConsumerWithRecordsToPoll(groupProtocol, tp, recordCount, new StringDeserializer());
    }

    private KafkaConsumer<?, ?> setUpConsumerWithRecordsToPoll(GroupProtocol groupProtocol,
                                                               TopicPartition tp,
                                                               int recordCount,
                                                               Deserializer<String> deserializer) {
        Cluster cluster = TestUtils.singletonCluster(tp.topic(), 1);
        Node node = cluster.nodes().get(0);

        ConsumerMetadata metadata = createMetadata(subscription);
        MockClient client = new MockClient(time, metadata);
        initMetadata(client, Map.of(topic, 1));
        consumer = newConsumer(groupProtocol, time, client, subscription, metadata, assignor,
                true, groupId, groupInstanceId, Optional.of(deserializer), false);
        consumer.subscribe(Set.of(topic), getConsumerRebalanceListener(consumer));
        prepareRebalance(client, node, assignor, List.of(tp), null);
        consumer.updateAssignmentMetadataIfNeeded(time.timer(Long.MAX_VALUE));
        client.prepareResponseFrom(fetchResponse(tp, 0, recordCount), node);
        return consumer;
    }

    @ParameterizedTest
    @EnumSource(GroupProtocol.class)
    public void testConstructorClose(GroupProtocol groupProtocol) {
        Properties props = new Properties();
        props.setProperty(ConsumerConfig.GROUP_PROTOCOL_CONFIG, groupProtocol.name());
        props.setProperty(ConsumerConfig.CLIENT_ID_CONFIG, "testConstructorClose");
        props.setProperty(ConsumerConfig.BOOTSTRAP_SERVERS_CONFIG, "invalid-23-8409-adsfsdj");
        props.setProperty(ConsumerConfig.METRIC_REPORTER_CLASSES_CONFIG, MockMetricsReporter.class.getName());

        final int oldInitCount = MockMetricsReporter.INIT_COUNT.get();
        final int oldCloseCount = MockMetricsReporter.CLOSE_COUNT.get();
        try {
            newConsumer(props, new ByteArrayDeserializer(), new ByteArrayDeserializer());
            fail("should have caught an exception and returned");
        } catch (KafkaException e) {
            assertEquals(oldInitCount + 1, MockMetricsReporter.INIT_COUNT.get());
            assertEquals(oldCloseCount + 1, MockMetricsReporter.CLOSE_COUNT.get());
            assertEquals("Failed to construct kafka consumer", e.getMessage());
        }
    }

    @ParameterizedTest
    @EnumSource(GroupProtocol.class)
    public void testConstructorInvalidMetricReporters(GroupProtocol groupProtocol) {
        Properties props = new Properties();
        props.setProperty(ConsumerConfig.GROUP_PROTOCOL_CONFIG, groupProtocol.name());
        props.setProperty(ConsumerConfig.CLIENT_ID_CONFIG, "testConstructorInvalidMetricReporters");
        props.setProperty(ConsumerConfig.BOOTSTRAP_SERVERS_CONFIG, "localhost:9999");
        props.setProperty(ConsumerConfig.METRIC_REPORTER_CLASSES_CONFIG, "an.invalid.class");

        KafkaException e = assertThrows(
                KafkaException.class,
                () -> newConsumer(props, new ByteArrayDeserializer(), new ByteArrayDeserializer()));

        assertEquals("Failed to construct kafka consumer", e.getMessage());
        assertEquals("Class an.invalid.class cannot be found", e.getCause().getMessage());
    }

    @ParameterizedTest
    @EnumSource(GroupProtocol.class)
    public void testOsDefaultSocketBufferSizes(GroupProtocol groupProtocol) {
        Map<String, Object> config = new HashMap<>();
        config.put(ConsumerConfig.GROUP_PROTOCOL_CONFIG, groupProtocol.name());
        config.put(ConsumerConfig.BOOTSTRAP_SERVERS_CONFIG, "localhost:9999");
        config.put(ConsumerConfig.SEND_BUFFER_CONFIG, Selectable.USE_DEFAULT_BUFFER_SIZE);
        config.put(ConsumerConfig.RECEIVE_BUFFER_CONFIG, Selectable.USE_DEFAULT_BUFFER_SIZE);
        consumer = newConsumer(config, new ByteArrayDeserializer(), new ByteArrayDeserializer());
    }

    @ParameterizedTest
    @EnumSource(GroupProtocol.class)
    public void testInvalidSocketSendBufferSize(GroupProtocol groupProtocol) {
        Map<String, Object> config = new HashMap<>();
        config.put(ConsumerConfig.GROUP_PROTOCOL_CONFIG, groupProtocol.name());
        config.put(ConsumerConfig.BOOTSTRAP_SERVERS_CONFIG, "localhost:9999");
        config.put(ConsumerConfig.SEND_BUFFER_CONFIG, -2);
        assertThrows(KafkaException.class,
            () -> newConsumer(config, new ByteArrayDeserializer(), new ByteArrayDeserializer()));
    }

    @ParameterizedTest
    @EnumSource(GroupProtocol.class)
    public void testInvalidSocketReceiveBufferSize(GroupProtocol groupProtocol) {
        Map<String, Object> config = new HashMap<>();
        config.put(ConsumerConfig.GROUP_PROTOCOL_CONFIG, groupProtocol.name());
        config.put(ConsumerConfig.BOOTSTRAP_SERVERS_CONFIG, "localhost:9999");
        config.put(ConsumerConfig.RECEIVE_BUFFER_CONFIG, -2);
        assertThrows(KafkaException.class,
            () -> newConsumer(config, new ByteArrayDeserializer(), new ByteArrayDeserializer()));
    }

    @ParameterizedTest
    @EnumSource(GroupProtocol.class)
    public void shouldIgnoreGroupInstanceIdForEmptyGroupId(GroupProtocol groupProtocol) {
        Map<String, Object> config = new HashMap<>();
        config.put(ConsumerConfig.GROUP_PROTOCOL_CONFIG, groupProtocol.name());
        config.put(ConsumerConfig.BOOTSTRAP_SERVERS_CONFIG, "localhost:9999");
        config.put(ConsumerConfig.GROUP_INSTANCE_ID_CONFIG, "instance_id");
        consumer = newConsumer(config, new ByteArrayDeserializer(), new ByteArrayDeserializer());
    }

    @ParameterizedTest
    @EnumSource(GroupProtocol.class)
    public void testSubscription(GroupProtocol groupProtocol) {
        consumer = newConsumer(groupProtocol, groupId);

        consumer.subscribe(List.of(topic));
        assertEquals(Set.of(topic), consumer.subscription());
        assertTrue(consumer.assignment().isEmpty());

        consumer.subscribe(Collections.emptyList());
        assertTrue(consumer.subscription().isEmpty());
        assertTrue(consumer.assignment().isEmpty());

        consumer.assign(List.of(tp0));
        assertTrue(consumer.subscription().isEmpty());
        assertEquals(Set.of(tp0), consumer.assignment());

        consumer.unsubscribe();
        assertTrue(consumer.subscription().isEmpty());
        assertTrue(consumer.assignment().isEmpty());
    }

    @ParameterizedTest
    @EnumSource(GroupProtocol.class)
    public void testSubscriptionOnNullTopicCollection(GroupProtocol groupProtocol) {
        consumer = newConsumer(groupProtocol, groupId);
        assertThrows(IllegalArgumentException.class, () -> consumer.subscribe((List<String>) null));
    }

    @ParameterizedTest
    @EnumSource(GroupProtocol.class)
    public void testSubscriptionOnNullTopic(GroupProtocol groupProtocol) {
        consumer = newConsumer(groupProtocol, groupId);
        assertThrows(IllegalArgumentException.class, () -> consumer.subscribe(singletonList(null)));
    }

    @ParameterizedTest
    @EnumSource(GroupProtocol.class)
    public void testSubscriptionOnEmptyTopic(GroupProtocol groupProtocol) {
        consumer = newConsumer(groupProtocol, groupId);
        String emptyTopic = "  ";
        assertThrows(IllegalArgumentException.class, () -> consumer.subscribe(List.of(emptyTopic)));
    }

    @ParameterizedTest
    @EnumSource(GroupProtocol.class)
    public void testSubscriptionOnNullPattern(GroupProtocol groupProtocol) {
        consumer = newConsumer(groupProtocol, groupId);
        assertThrows(IllegalArgumentException.class,
            () -> consumer.subscribe((Pattern) null));
    }

    @ParameterizedTest
    @EnumSource(GroupProtocol.class)
    public void testSubscriptionOnEmptyPattern(GroupProtocol groupProtocol) {
        consumer = newConsumer(groupProtocol, groupId);
        assertThrows(IllegalArgumentException.class,
            () -> consumer.subscribe(Pattern.compile("")));
    }

    @ParameterizedTest
    @EnumSource(value = GroupProtocol.class, names = "CLASSIC")
    public void testSubscriptionWithEmptyPartitionAssignment(GroupProtocol groupProtocol) {
        Properties props = new Properties();
        props.setProperty(ConsumerConfig.GROUP_PROTOCOL_CONFIG, groupProtocol.name());
        props.setProperty(ConsumerConfig.BOOTSTRAP_SERVERS_CONFIG, "localhost:9999");
        props.setProperty(ConsumerConfig.PARTITION_ASSIGNMENT_STRATEGY_CONFIG, "");
        props.setProperty(ConsumerConfig.GROUP_ID_CONFIG, groupId);
        consumer = newConsumer(props, new ByteArrayDeserializer(), new ByteArrayDeserializer());
        assertThrows(IllegalStateException.class,
            () -> consumer.subscribe(List.of(topic)));
    }

    @ParameterizedTest
    @EnumSource(GroupProtocol.class)
    public void testSeekNegative(GroupProtocol groupProtocol) {
        consumer = newConsumer(groupProtocol, null);
        consumer.assign(Set.of(new TopicPartition("nonExistTopic", 0)));
        assertThrows(IllegalArgumentException.class,
            () -> consumer.seek(new TopicPartition("nonExistTopic", 0), -1));
    }

    @ParameterizedTest
    @EnumSource(GroupProtocol.class)
    public void testAssignOnNullTopicPartition(GroupProtocol groupProtocol) {
        consumer = newConsumer(groupProtocol, null);
        assertThrows(IllegalArgumentException.class, () -> consumer.assign(null));
    }

    @ParameterizedTest
    @EnumSource(GroupProtocol.class)
    public void testAssignOnEmptyTopicPartition(GroupProtocol groupProtocol) {
        consumer = newConsumer(groupProtocol, groupId);
        consumer.assign(Collections.emptyList());
        assertTrue(consumer.subscription().isEmpty());
        assertTrue(consumer.assignment().isEmpty());
    }

    @ParameterizedTest
    @EnumSource(GroupProtocol.class)
    public void testAssignOnNullTopicInPartition(GroupProtocol groupProtocol) {
        consumer = newConsumer(groupProtocol, null);
        assertThrows(IllegalArgumentException.class, () -> consumer.assign(Set.of(new TopicPartition(null, 0))));
    }

    @ParameterizedTest
    @EnumSource(GroupProtocol.class)
    public void testAssignOnEmptyTopicInPartition(GroupProtocol groupProtocol) {
        consumer = newConsumer(groupProtocol, null);
        assertThrows(IllegalArgumentException.class, () -> consumer.assign(Set.of(new TopicPartition("  ", 0))));
    }

    @ParameterizedTest
    @EnumSource(GroupProtocol.class)
    public void testInterceptorConstructorClose(GroupProtocol groupProtocol) {
        try {
            Properties props = new Properties();
            // test with client ID assigned by KafkaConsumer
            props.setProperty(ConsumerConfig.GROUP_PROTOCOL_CONFIG, groupProtocol.name());
            props.setProperty(ConsumerConfig.BOOTSTRAP_SERVERS_CONFIG, "localhost:9999");
            props.setProperty(ConsumerConfig.INTERCEPTOR_CLASSES_CONFIG, MockConsumerInterceptor.class.getName());

            consumer = newConsumer(
                    props, new StringDeserializer(), new StringDeserializer());
            assertEquals(1, MockConsumerInterceptor.INIT_COUNT.get());
            assertEquals(0, MockConsumerInterceptor.CLOSE_COUNT.get());

            consumer.close(CloseOptions.timeout(Duration.ZERO));
            assertEquals(1, MockConsumerInterceptor.INIT_COUNT.get());
            assertEquals(1, MockConsumerInterceptor.CLOSE_COUNT.get());
            // Cluster metadata will only be updated on calling poll.
            assertNull(MockConsumerInterceptor.CLUSTER_META.get());

        } finally {
            // cleanup since we are using mutable static variables in MockConsumerInterceptor
            MockConsumerInterceptor.resetCounters();
        }
    }

    @ParameterizedTest
    @EnumSource(GroupProtocol.class)
    public void testInterceptorConstructorConfigurationWithExceptionShouldCloseRemainingInstances(GroupProtocol groupProtocol) {
        final int targetInterceptor = 1;

        try {
            Properties props = new Properties();
            props.setProperty(ConsumerConfig.GROUP_PROTOCOL_CONFIG, groupProtocol.name());
            props.setProperty(ConsumerConfig.BOOTSTRAP_SERVERS_CONFIG, "localhost:9999");
            props.setProperty(ConsumerConfig.INTERCEPTOR_CLASSES_CONFIG, 
                    CloseInterceptor.class.getName() + "," + MockConsumerInterceptor.class.getName());

            MockConsumerInterceptor.setThrowOnConfigExceptionThreshold(targetInterceptor);

            assertThrows(KafkaException.class, () -> newConsumer(
                    props, new StringDeserializer(), new StringDeserializer()));

            assertEquals(1, MockConsumerInterceptor.CONFIG_COUNT.get());
            assertEquals(1, MockConsumerInterceptor.CLOSE_COUNT.get());

            assertEquals(1, CloseInterceptor.CLOSE_COUNT.get());
        } finally {
            MockConsumerInterceptor.resetCounters();
            CloseInterceptor.resetCounters();
        }
    }

    @ParameterizedTest
    @EnumSource(GroupProtocol.class)
    public void testPause(GroupProtocol groupProtocol) {
        consumer = newConsumer(groupProtocol, groupId);

        consumer.assign(List.of(tp0));
        assertEquals(Set.of(tp0), consumer.assignment());
        assertTrue(consumer.paused().isEmpty());

        consumer.pause(Set.of(tp0));
        assertEquals(Set.of(tp0), consumer.paused());

        consumer.resume(Set.of(tp0));
        assertTrue(consumer.paused().isEmpty());

        consumer.unsubscribe();
        assertTrue(consumer.paused().isEmpty());
    }

    @ParameterizedTest
    @EnumSource(GroupProtocol.class)
    public void testConsumerJmxPrefix(GroupProtocol groupProtocol) throws  Exception {
        Map<String, Object> config = new HashMap<>();
        config.put(ConsumerConfig.GROUP_PROTOCOL_CONFIG, groupProtocol.name());
        config.put(ConsumerConfig.BOOTSTRAP_SERVERS_CONFIG, "localhost:9999");
        config.put(ConsumerConfig.SEND_BUFFER_CONFIG, Selectable.USE_DEFAULT_BUFFER_SIZE);
        config.put(ConsumerConfig.RECEIVE_BUFFER_CONFIG, Selectable.USE_DEFAULT_BUFFER_SIZE);
        config.put("client.id", "client-1");
        consumer = newConsumer(config, new ByteArrayDeserializer(), new ByteArrayDeserializer());
        MBeanServer server = ManagementFactory.getPlatformMBeanServer();
        MetricName testMetricName = consumer.metricsRegistry().metricName("test-metric",
                "grp1", "test metric");
        consumer.metricsRegistry().addMetric(testMetricName, new Avg());
        assertNotNull(server.getObjectInstance(new ObjectName("kafka.consumer:type=grp1,client-id=client-1")));
    }

    private KafkaConsumer<byte[], byte[]> newConsumer(GroupProtocol groupProtocol, String groupId) {
        return newConsumer(groupProtocol, groupId, Optional.empty());
    }

    private KafkaConsumer<byte[], byte[]> newConsumer(GroupProtocol groupProtocol,
                                                      String groupId,
                                                      Optional<Boolean> enableAutoCommit) {
        Properties props = new Properties();
        props.setProperty(ConsumerConfig.GROUP_PROTOCOL_CONFIG, groupProtocol.name());
        props.setProperty(ConsumerConfig.CLIENT_ID_CONFIG, "my.consumer");
        props.setProperty(ConsumerConfig.BOOTSTRAP_SERVERS_CONFIG, "localhost:9999");
        props.setProperty(ConsumerConfig.METRIC_REPORTER_CLASSES_CONFIG, MockMetricsReporter.class.getName());
        if (groupId != null)
            props.setProperty(ConsumerConfig.GROUP_ID_CONFIG, groupId);
        enableAutoCommit.ifPresent(
            autoCommit -> props.setProperty(ConsumerConfig.ENABLE_AUTO_COMMIT_CONFIG, autoCommit.toString()));
        return newConsumer(props, new ByteArrayDeserializer(), new ByteArrayDeserializer());
    }

    private <K, V> KafkaConsumer<K, V> newConsumer(Properties props) {
        return newConsumer(props, null, null);
    }

    private <K, V> KafkaConsumer<K, V> newConsumer(Map<String, Object> configs,
                                                   Deserializer<K> keyDeserializer,
                                                   Deserializer<V> valueDeserializer) {
        return new KafkaConsumer<>(new ConsumerConfig(ConsumerConfig.appendDeserializerToConfig(configs, keyDeserializer, valueDeserializer)),
                keyDeserializer, valueDeserializer);
    }

    private <K, V> KafkaConsumer<K, V> newConsumer(Properties props,
                                                   Deserializer<K> keyDeserializer,
                                                   Deserializer<V> valueDeserializer) {
        return newConsumer(propsToMap(props), keyDeserializer, valueDeserializer);
    }

    // TODO: this test requires rebalance logic which is not yet implemented in the CONSUMER group protocol.
    //       Once it is implemented, this should use both group protocols.
    @ParameterizedTest
    @EnumSource(value = GroupProtocol.class, names = "CLASSIC")
    public void verifyHeartbeatSent(GroupProtocol groupProtocol) throws Exception {
        ConsumerMetadata metadata = createMetadata(subscription);
        MockClient client = new MockClient(time, metadata);

        initMetadata(client, Map.of(topic, 1));
        Node node = metadata.fetch().nodes().get(0);

        consumer = newConsumer(groupProtocol, time, client, subscription, metadata, assignor, true, groupInstanceId);

        consumer.subscribe(Set.of(topic), getConsumerRebalanceListener(consumer));
        Node coordinator = prepareRebalance(client, node, assignor, List.of(tp0), null);

        // initial fetch
        client.prepareResponseFrom(fetchResponse(tp0, 0, 0), node);
        consumer.updateAssignmentMetadataIfNeeded(time.timer(Long.MAX_VALUE));

        assertEquals(Set.of(tp0), consumer.assignment());

        AtomicBoolean heartbeatReceived = prepareHeartbeatResponse(client, coordinator, Errors.NONE);

        // heartbeat interval is 2 seconds
        time.sleep(heartbeatIntervalMs);
        Thread.sleep(heartbeatIntervalMs);

        consumer.updateAssignmentMetadataIfNeeded(time.timer(Long.MAX_VALUE));

        assertTrue(heartbeatReceived.get());
    }

    // TODO: this test requires rebalance logic which is not yet implemented in the CONSUMER group protocol.
    //       Once it is implemented, this should use both group protocols.
    @ParameterizedTest
    @EnumSource(value = GroupProtocol.class, names = "CLASSIC")
    public void verifyHeartbeatSentWhenFetchedDataReady(GroupProtocol groupProtocol) throws Exception {
        ConsumerMetadata metadata = createMetadata(subscription);
        MockClient client = new MockClient(time, metadata);

        initMetadata(client, Map.of(topic, 1));
        Node node = metadata.fetch().nodes().get(0);

        consumer = newConsumer(groupProtocol, time, client, subscription, metadata, assignor, true, groupInstanceId);
        consumer.subscribe(Set.of(topic), getConsumerRebalanceListener(consumer));
        Node coordinator = prepareRebalance(client, node, assignor, List.of(tp0), null);

        consumer.updateAssignmentMetadataIfNeeded(time.timer(Long.MAX_VALUE));
        consumer.poll(Duration.ZERO);

        // respond to the outstanding fetch so that we have data available on the next poll
        client.respondFrom(fetchResponse(tp0, 0, 5), node);
        client.poll(0, time.milliseconds());

        client.prepareResponseFrom(fetchResponse(tp0, 5, 0), node);
        AtomicBoolean heartbeatReceived = prepareHeartbeatResponse(client, coordinator, Errors.NONE);

        time.sleep(heartbeatIntervalMs);
        Thread.sleep(heartbeatIntervalMs);

        consumer.poll(Duration.ZERO);

        assertTrue(heartbeatReceived.get());
    }

    @ParameterizedTest
    @EnumSource(GroupProtocol.class)
    public void verifyPollTimesOutDuringMetadataUpdate(GroupProtocol groupProtocol) {
        final ConsumerMetadata metadata = createMetadata(subscription);
        final MockClient client = new MockClient(time, metadata);

        initMetadata(client, Map.of(topic, 1));
        Node node = metadata.fetch().nodes().get(0);

        client.prepareResponseFrom(FindCoordinatorResponse.prepareResponse(Errors.NONE, groupId, node), node);
        consumer = newConsumer(groupProtocol, time, client, subscription, metadata, assignor, true, groupInstanceId);
        consumer.subscribe(Set.of(topic), getConsumerRebalanceListener(consumer));
        // Since we would enable the heartbeat thread after received join-response which could
        // send the sync-group on behalf of the consumer if it is enqueued, we may still complete
        // the rebalance and send out the fetch; in order to avoid it we do not prepare sync response here.
        Node coordinator = new Node(Integer.MAX_VALUE - node.id(), node.host(), node.port());
        client.prepareResponseFrom(joinGroupFollowerResponse(assignor, 1, memberId, leaderId, Errors.NONE), coordinator);

        consumer.poll(Duration.ZERO);

        final Queue<ClientRequest> requests = client.requests();
        assertEquals(0, requests.stream().filter(request -> request.apiKey().equals(ApiKeys.FETCH)).count());
    }

    @ParameterizedTest
    @EnumSource(GroupProtocol.class)
    @SuppressWarnings("unchecked")
    public void verifyNoCoordinatorLookupForManualAssignmentWithSeek(GroupProtocol groupProtocol) {
        ConsumerMetadata metadata = createMetadata(subscription);
        MockClient client = new MockClient(time, metadata);

        initMetadata(client, Map.of(topic, 1));

        consumer = newConsumer(groupProtocol, time, client, subscription, metadata, assignor, false, null, groupInstanceId, false);
        consumer.assign(Set.of(tp0));
        consumer.seekToBeginning(Set.of(tp0));

        // there shouldn't be any need to look up the coordinator or fetch committed offsets.
        // we just look up the starting position and send the record fetch.
        client.prepareResponse(listOffsetsResponse(Map.of(tp0, 50L)));
        client.prepareResponse(fetchResponse(tp0, 50L, 5));

        ConsumerRecords<String, String> records = (ConsumerRecords<String, String>) consumer.poll(Duration.ofMillis(1));
        assertEquals(5, records.count());
        assertEquals(55L, consumer.position(tp0));
        assertEquals(1, records.nextOffsets().size());
        assertEquals(new OffsetAndMetadata(55L), records.nextOffsets().get(tp0));
    }

    // TODO: this test triggers a bug with the CONSUMER group protocol implementation.
    //       The bug will be investigated and fixed so this test can use both group protocols.
    @ParameterizedTest
    @EnumSource(value = GroupProtocol.class, names = "CLASSIC")
    public void verifyNoCoordinatorLookupForManualAssignmentWithOffsetCommit(GroupProtocol groupProtocol) {
        ConsumerMetadata metadata = createMetadata(subscription);
        MockClient client = new MockClient(time, metadata);

        initMetadata(client, Map.of(topic, 1));
        Node node = metadata.fetch().nodes().get(0);

        // create a consumer with groupID with manual assignment
        consumer = newConsumer(groupProtocol, time, client, subscription, metadata, assignor, true, groupInstanceId);
        consumer.assign(Set.of(tp0));

        // 1st coordinator error should cause coordinator unknown
        client.prepareResponseFrom(FindCoordinatorResponse.prepareResponse(Errors.COORDINATOR_NOT_AVAILABLE, groupId, node), node);
        consumer.poll(Duration.ofMillis(0));

        // 2nd coordinator error should find the correct coordinator and clear the findCoordinatorFuture
        client.prepareResponseFrom(FindCoordinatorResponse.prepareResponse(Errors.NONE, groupId, node), node);

        client.prepareResponse(offsetResponse(Map.of(tp0, 50L), Errors.NONE));
        client.prepareResponse(fetchResponse(tp0, 50L, 5));

        @SuppressWarnings("unchecked")
        ConsumerRecords<String, String> records = (ConsumerRecords<String, String>) consumer.poll(Duration.ofMillis(0));
        assertEquals(5, records.count());
        assertEquals(55L, consumer.position(tp0));
        assertEquals(1, records.nextOffsets().size());
        assertEquals(new OffsetAndMetadata(55L), records.nextOffsets().get(tp0));

        // after coordinator found, consumer should be able to commit the offset successfully
        client.prepareResponse(offsetCommitResponse(Map.of(tp0, Errors.NONE)));
        consumer.commitSync(Map.of(tp0, new OffsetAndMetadata(55L)));

        // verify the offset is committed
        client.prepareResponse(offsetResponse(Map.of(tp0, 55L), Errors.NONE));
        assertEquals(55, consumer.committed(Set.of(tp0), Duration.ZERO).get(tp0).offset());
    }

    @ParameterizedTest
    @EnumSource(value = GroupProtocol.class)
    public void testFetchProgressWithMissingPartitionPosition(GroupProtocol groupProtocol) {
        // Verifies that we can make progress on one partition while we are awaiting
        // a reset on another partition.

        ConsumerMetadata metadata = createMetadata(subscription);
        MockClient client = new MockClient(time, metadata);
        initMetadata(client, Map.of(topic, 2));

        if (groupProtocol == GroupProtocol.CONSUMER) {
            Node node = metadata.fetch().nodes().get(0);
            client.prepareResponseFrom(FindCoordinatorResponse.prepareResponse(Errors.NONE, groupId, node), node);
        }

        consumer = newConsumerNoAutoCommit(groupProtocol, time, client, subscription, metadata);
        consumer.assign(Arrays.asList(tp0, tp1));
        consumer.seekToEnd(Set.of(tp0));
        consumer.seekToBeginning(Set.of(tp1));

        client.prepareResponse(body -> {
            ListOffsetsRequest request = (ListOffsetsRequest) body;
            List<ListOffsetsPartition> partitions = request.topics().stream().flatMap(t -> {
                if (t.name().equals(topic))
                    return Stream.of(t.partitions());
                else
                    return Stream.empty();
            }).flatMap(List::stream).collect(Collectors.toList());
            ListOffsetsPartition expectedTp0 = new ListOffsetsPartition()
                    .setPartitionIndex(tp0.partition())
                    .setTimestamp(ListOffsetsRequest.LATEST_TIMESTAMP);
            ListOffsetsPartition expectedTp1 = new ListOffsetsPartition()
                    .setPartitionIndex(tp1.partition())
                    .setTimestamp(ListOffsetsRequest.EARLIEST_TIMESTAMP);
            return partitions.contains(expectedTp0) && partitions.contains(expectedTp1);
        }, listOffsetsResponse(Map.of(tp0, 50L), Map.of(tp1, Errors.NOT_LEADER_OR_FOLLOWER)));
        client.prepareResponse(
            body -> {
                FetchRequest request = (FetchRequest) body;
                Map<TopicIdPartition, FetchRequest.PartitionData> fetchData = request.fetchData(topicNames);
                TopicIdPartition tidp0 = new TopicIdPartition(topicIds.get(tp0.topic()), tp0);
                return fetchData.keySet().equals(Set.of(tidp0)) &&
                        fetchData.get(tidp0).fetchOffset == 50L;

            }, fetchResponse(tp0, 50L, 5));

        @SuppressWarnings("unchecked")
        ConsumerRecords<String, String> records = (ConsumerRecords<String, String>) consumer.poll(Duration.ofMillis(1));
        assertEquals(5, records.count());
        assertEquals(Set.of(tp0), records.partitions());
        assertEquals(1, records.nextOffsets().size());
        assertEquals(new OffsetAndMetadata(records.records(tp0).get(records.count() - 1).offset() + 1), records.nextOffsets().get(tp0));
    }

    private void initMetadata(MockClient mockClient, Map<String, Integer> partitionCounts) {
        Map<String, Uuid> metadataIds = new HashMap<>();
        for (String name : partitionCounts.keySet()) {
            metadataIds.put(name, topicIds.get(name));
        }
        MetadataResponse initialMetadata = RequestTestUtils.metadataUpdateWithIds(1, partitionCounts, metadataIds);

        mockClient.updateMetadata(initialMetadata);
    }

    @ParameterizedTest
    @EnumSource(value = GroupProtocol.class)
    public void testMissingOffsetNoResetPolicy(GroupProtocol groupProtocol) throws InterruptedException {
        SubscriptionState subscription = new SubscriptionState(new LogContext(), AutoOffsetResetStrategy.NONE);
        ConsumerMetadata metadata = createMetadata(subscription);
        MockClient client = new MockClient(time, metadata);

        initMetadata(client, Map.of(topic, 1));
        Node node = metadata.fetch().nodes().get(0);

        client.prepareResponseFrom(FindCoordinatorResponse.prepareResponse(Errors.NONE, groupId, node), node);
        Node coordinator = new Node(Integer.MAX_VALUE - node.id(), node.host(), node.port());
        client.prepareResponseFrom(offsetResponse(Map.of(tp0, -1L), Errors.NONE), coordinator);

        consumer = newConsumer(groupProtocol, time, client, subscription, metadata, assignor,
                true, groupId, groupInstanceId, false);
        consumer.assign(List.of(tp0));

        if (groupProtocol == GroupProtocol.CONSUMER) {
            // New consumer poll(ZERO) needs to wait for the offset fetch event added by a call to poll, to be processed
            // by the background thread, so it can realize there are no committed offsets and then
            // throw the NoOffsetForPartitionException
            assertPollEventuallyThrows(consumer, NoOffsetForPartitionException.class,
                    "Consumer was not able to update fetch positions on continuous calls with 0 timeout");
        } else {
            assertThrows(NoOffsetForPartitionException.class, () -> consumer.poll(Duration.ZERO));
        }
    }

    @ParameterizedTest
    @EnumSource(GroupProtocol.class)
    public void testResetToCommittedOffset(GroupProtocol groupProtocol) {
        SubscriptionState subscription = new SubscriptionState(new LogContext(), AutoOffsetResetStrategy.NONE);
        ConsumerMetadata metadata = createMetadata(subscription);
        MockClient client = new MockClient(time, metadata);

        initMetadata(client, Map.of(topic, 1));
        Node node = metadata.fetch().nodes().get(0);

        client.prepareResponseFrom(FindCoordinatorResponse.prepareResponse(Errors.NONE, groupId, node), node);
        KafkaConsumer<String, String> consumer = newConsumer(groupProtocol, time, client, subscription, metadata, assignor,
                true, groupId, groupInstanceId, false);
        consumer.assign(List.of(tp0));

        Node coordinator = new Node(Integer.MAX_VALUE - node.id(), node.host(), node.port());

        client.prepareResponseFrom(offsetResponse(Map.of(tp0, 539L), Errors.NONE), coordinator);
        consumer.poll(Duration.ZERO);

        assertEquals(539L, consumer.position(tp0));
    }

    @ParameterizedTest
    @EnumSource(GroupProtocol.class)
    public void testResetUsingAutoResetPolicy(GroupProtocol groupProtocol) {
        setUpConsumerWithAutoResetPolicy(groupProtocol, AutoOffsetResetStrategy.LATEST);
        assertEquals(50L, consumer.position(tp0));
    }

    @ParameterizedTest
    @EnumSource(GroupProtocol.class)
    public void testResetUsingDurationBasedAutoResetPolicy(GroupProtocol groupProtocol) {
        AutoOffsetResetStrategy durationStrategy = AutoOffsetResetStrategy.fromString("by_duration:PT1H");
        setUpConsumerWithAutoResetPolicy(groupProtocol, durationStrategy);
        assertEquals(50L, consumer.position(tp0));
    }

    private void setUpConsumerWithAutoResetPolicy(GroupProtocol groupProtocol, AutoOffsetResetStrategy strategy) {
        SubscriptionState subscription = new SubscriptionState(new LogContext(), strategy);
        ConsumerMetadata metadata = createMetadata(subscription);
        MockClient client = new MockClient(time, metadata);

        initMetadata(client, Map.of(topic, 1));
        Node node = metadata.fetch().nodes().get(0);

        client.prepareResponseFrom(FindCoordinatorResponse.prepareResponse(Errors.NONE, groupId, node), node);
        consumer = newConsumer(groupProtocol, time, client, subscription, metadata, assignor,
                true, groupId, groupInstanceId, false);
        consumer.assign(List.of(tp0));

        Node coordinator = new Node(Integer.MAX_VALUE - node.id(), node.host(), node.port());

        client.prepareResponseFrom(offsetResponse(Map.of(tp0, -1L), Errors.NONE), coordinator);
        client.prepareResponse(listOffsetsResponse(Map.of(tp0, 50L)));

        consumer.poll(Duration.ZERO);
    }

    @ParameterizedTest
    @EnumSource(GroupProtocol.class)
    public void testOffsetIsValidAfterSeek(GroupProtocol groupProtocol) {
        SubscriptionState subscription = new SubscriptionState(new LogContext(), AutoOffsetResetStrategy.LATEST);
        ConsumerMetadata metadata = createMetadata(subscription);
        MockClient client = new MockClient(time, metadata);

        initMetadata(client, Map.of(topic, 1));

        consumer = newConsumer(groupProtocol, time, client, subscription, metadata, assignor,
                true, groupId, Optional.empty(), false);
        consumer.assign(List.of(tp0));
        consumer.seek(tp0, 20L);
        consumer.poll(Duration.ZERO);
        assertEquals(20L, subscription.validPosition(tp0).offset);
    }

    @ParameterizedTest
    @EnumSource(GroupProtocol.class)
    public void testCommitsFetchedDuringAssign(GroupProtocol groupProtocol) {
        long offset1 = 10000;
        long offset2 = 20000;

        ConsumerMetadata metadata = createMetadata(subscription);
        MockClient client = new MockClient(time, metadata);

        initMetadata(client, Map.of(topic, 2));
        Node node = metadata.fetch().nodes().get(0);

        client.prepareResponseFrom(FindCoordinatorResponse.prepareResponse(Errors.NONE, groupId, node), node);
        consumer = newConsumer(groupProtocol, time, client, subscription, metadata, assignor, true, groupInstanceId);
        consumer.assign(List.of(tp0));

        // lookup coordinator
        Node coordinator = new Node(Integer.MAX_VALUE - node.id(), node.host(), node.port());

        // fetch offset for one topic
        client.prepareResponseFrom(offsetResponse(Map.of(tp0, offset1), Errors.NONE), coordinator);
        assertEquals(offset1, consumer.committed(Set.of(tp0)).get(tp0).offset());

        consumer.assign(Arrays.asList(tp0, tp1));

        // fetch offset for two topics
        Map<TopicPartition, Long> offsets = new HashMap<>();
        offsets.put(tp0, offset1);
        client.prepareResponseFrom(offsetResponse(offsets, Errors.NONE), coordinator);
        assertEquals(offset1, consumer.committed(Set.of(tp0)).get(tp0).offset());

        offsets.remove(tp0);
        offsets.put(tp1, offset2);
        client.prepareResponseFrom(offsetResponse(offsets, Errors.NONE), coordinator);
        assertEquals(offset2, consumer.committed(Set.of(tp1)).get(tp1).offset());
    }

    @SuppressWarnings("resource")
    @ParameterizedTest
    @EnumSource(GroupProtocol.class)
    public void testFetchStableOffsetThrowInCommitted(GroupProtocol groupProtocol) {
        assertThrows(UnsupportedVersionException.class, () -> setupThrowableConsumer(groupProtocol).committed(Set.of(tp0)));
    }

    @SuppressWarnings("resource")
    @ParameterizedTest
    @EnumSource(GroupProtocol.class)
    public void testFetchStableOffsetThrowInPoll(GroupProtocol groupProtocol) throws InterruptedException {
        setupThrowableConsumer(groupProtocol);
        TestUtils.waitForCondition(() -> {
            try {
                // For CONSUMER protocol, the offset request is sent in the background thread,
                // so we need to wait until the offset request is sent.
                consumer.poll(Duration.ZERO);
                return false;
            } catch (UnsupportedVersionException e) {
                return true;
            }
        }, "Failed to throw UnsupportedVersionException in poll");
    }

    @SuppressWarnings("resource")
    @ParameterizedTest
    @EnumSource(GroupProtocol.class)
    public void testFetchStableOffsetThrowInPosition(GroupProtocol groupProtocol) {
        assertThrows(UnsupportedVersionException.class, () -> setupThrowableConsumer(groupProtocol).position(tp0));
    }

    private KafkaConsumer<?, ?> setupThrowableConsumer(GroupProtocol groupProtocol) {
        long offset1 = 10000;

        ConsumerMetadata metadata = createMetadata(subscription);
        MockClient client = new MockClient(time, metadata);

        initMetadata(client, Map.of(topic, 2));
        client.setNodeApiVersions(NodeApiVersions.create(ApiKeys.OFFSET_FETCH.id, (short) 0, (short) 6));

        Node node = metadata.fetch().nodes().get(0);

        client.prepareResponseFrom(FindCoordinatorResponse.prepareResponse(Errors.NONE, groupId, node), node);
        consumer = newConsumer(
                groupProtocol, time, client, subscription, metadata, assignor, true, groupId, groupInstanceId, true);
        consumer.assign(List.of(tp0));

        Node coordinator = new Node(Integer.MAX_VALUE - node.id(), node.host(), node.port());

        client.prepareResponseFrom(offsetResponse(
            Map.of(tp0, offset1), Errors.NONE), coordinator);
        return consumer;
    }

    @ParameterizedTest
    @EnumSource(GroupProtocol.class)
    public void testNoCommittedOffsets(GroupProtocol groupProtocol) {
        long offset1 = 10000;

        ConsumerMetadata metadata = createMetadata(subscription);
        MockClient client = new MockClient(time, metadata);

        initMetadata(client, Map.of(topic, 2));
        Node node = metadata.fetch().nodes().get(0);

        client.prepareResponseFrom(FindCoordinatorResponse.prepareResponse(Errors.NONE, groupId, node), node);
        consumer = newConsumer(groupProtocol, time, client, subscription, metadata, assignor, true, groupInstanceId);
        consumer.assign(Arrays.asList(tp0, tp1));

        // lookup coordinator
        Node coordinator = new Node(Integer.MAX_VALUE - node.id(), node.host(), node.port());

        // fetch offset for one topic
        client.prepareResponseFrom(offsetResponse(Utils.mkMap(Utils.mkEntry(tp0, offset1), Utils.mkEntry(tp1, -1L)), Errors.NONE), coordinator);
        final Map<TopicPartition, OffsetAndMetadata> committed = consumer.committed(Set.of(tp0, tp1));
        assertEquals(2, committed.size());
        assertEquals(offset1, committed.get(tp0).offset());
        assertNull(committed.get(tp1));
    }

    // TODO: this test requires rebalance logic which is not yet implemented in the CONSUMER group protocol.
    //       Once it is implemented, this should use both group protocols.
    @ParameterizedTest
    @EnumSource(value = GroupProtocol.class, names = "CLASSIC")
    public void testAutoCommitSentBeforePositionUpdate(GroupProtocol groupProtocol) {
        ConsumerMetadata metadata = createMetadata(subscription);
        MockClient client = new MockClient(time, metadata);

        initMetadata(client, Map.of(topic, 1));
        Node node = metadata.fetch().nodes().get(0);

        consumer = newConsumer(groupProtocol, time, client, subscription, metadata, assignor, true, groupInstanceId);
        consumer.subscribe(Set.of(topic), getConsumerRebalanceListener(consumer));
        Node coordinator = prepareRebalance(client, node, assignor, List.of(tp0), null);

        consumer.updateAssignmentMetadataIfNeeded(time.timer(Long.MAX_VALUE));
        consumer.poll(Duration.ZERO);

        // respond to the outstanding fetch so that we have data available on the next poll
        client.respondFrom(fetchResponse(tp0, 0, 5), node);
        client.poll(0, time.milliseconds());

        time.sleep(autoCommitIntervalMs);

        client.prepareResponseFrom(fetchResponse(tp0, 5, 0), node);

        // no data has been returned to the user yet, so the committed offset should be 0
        AtomicBoolean commitReceived = prepareOffsetCommitResponse(client, coordinator, tp0, 0);

        consumer.poll(Duration.ZERO);

        assertTrue(commitReceived.get());
    }

    // TODO: this test requires rebalance logic which is not yet implemented in the CONSUMER group protocol.
    //       Once it is implemented, this should use both group protocols.
    @ParameterizedTest
    @EnumSource(value = GroupProtocol.class, names = "CLASSIC")
    public void testRegexSubscription(GroupProtocol groupProtocol) {
        String unmatchedTopic = "unmatched";
        ConsumerMetadata metadata = createMetadata(subscription);
        MockClient client = new MockClient(time, metadata);

        Map<String, Integer> partitionCounts = new HashMap<>();
        partitionCounts.put(topic, 1);
        partitionCounts.put(unmatchedTopic, 1);
        topicIds.put(unmatchedTopic, Uuid.randomUuid());
        initMetadata(client, partitionCounts);
        Node node = metadata.fetch().nodes().get(0);

        consumer = newConsumer(groupProtocol, time, client, subscription, metadata, assignor, true, groupInstanceId);
        prepareRebalance(client, node, Set.of(topic), assignor, List.of(tp0), null);

        consumer.subscribe(Pattern.compile(topic), getConsumerRebalanceListener(consumer));

        client.prepareMetadataUpdate(RequestTestUtils.metadataUpdateWithIds(1, partitionCounts, topicIds));

        consumer.updateAssignmentMetadataIfNeeded(time.timer(Long.MAX_VALUE));

        assertEquals(Set.of(topic), consumer.subscription());
        assertEquals(Set.of(tp0), consumer.assignment());
    }

    // TODO: this test requires rebalance logic which is not yet implemented in the CONSUMER group protocol.
    //       Once it is implemented, this should use both group protocols.
    @ParameterizedTest
    @EnumSource(value = GroupProtocol.class, names = "CLASSIC")
    public void testChangingRegexSubscription(GroupProtocol groupProtocol) {
        String otherTopic = "other";
        TopicPartition otherTopicPartition = new TopicPartition(otherTopic, 0);

        ConsumerMetadata metadata = createMetadata(subscription);
        MockClient client = new MockClient(time, metadata);

        Map<String, Integer> partitionCounts = new HashMap<>();
        partitionCounts.put(topic, 1);
        partitionCounts.put(otherTopic, 1);
        topicIds.put(otherTopic, Uuid.randomUuid());
        initMetadata(client, partitionCounts);
        Node node = metadata.fetch().nodes().get(0);

        consumer = newConsumer(groupProtocol, time, client, subscription, metadata, assignor, false, groupInstanceId);

        Node coordinator = prepareRebalance(client, node, Set.of(topic), assignor, List.of(tp0), null);
        consumer.subscribe(Pattern.compile(topic), getConsumerRebalanceListener(consumer));

        consumer.updateAssignmentMetadataIfNeeded(time.timer(Long.MAX_VALUE));
        consumer.poll(Duration.ZERO);

        assertEquals(Set.of(topic), consumer.subscription());

        consumer.subscribe(Pattern.compile(otherTopic), getConsumerRebalanceListener(consumer));

        client.prepareMetadataUpdate(RequestTestUtils.metadataUpdateWithIds(1, partitionCounts, topicIds));
        prepareRebalance(client, node, Set.of(otherTopic), assignor, List.of(otherTopicPartition), coordinator);
        consumer.poll(Duration.ZERO);

        assertEquals(Set.of(otherTopic), consumer.subscription());
    }

    // TODO: this test requires rebalance logic which is not yet implemented in the CONSUMER group protocol.
    //       Once it is implemented, this should use both group protocols.
    @ParameterizedTest
    @EnumSource(value = GroupProtocol.class, names = "CLASSIC")
    public void testWakeupWithFetchDataAvailable(GroupProtocol groupProtocol) throws Exception {
        ConsumerMetadata metadata = createMetadata(subscription);
        MockClient client = new MockClient(time, metadata);

        initMetadata(client, Map.of(topic, 1));
        Node node = metadata.fetch().nodes().get(0);

        consumer = newConsumer(groupProtocol, time, client, subscription, metadata, assignor, true, groupInstanceId);
        consumer.subscribe(Set.of(topic), getConsumerRebalanceListener(consumer));
        prepareRebalance(client, node, assignor, List.of(tp0), null);

        consumer.updateAssignmentMetadataIfNeeded(time.timer(Long.MAX_VALUE));
        consumer.poll(Duration.ZERO);

        // respond to the outstanding fetch so that we have data available on the next poll
        client.respondFrom(fetchResponse(tp0, 0, 5), node);
        client.poll(0, time.milliseconds());

        consumer.wakeup();

        assertThrows(WakeupException.class, () -> consumer.poll(Duration.ZERO));

        // make sure the position hasn't been updated
        assertEquals(0, consumer.position(tp0));

        // the next poll should return the completed fetch
        @SuppressWarnings("unchecked")
        ConsumerRecords<String, String> records = (ConsumerRecords<String, String>) consumer.poll(Duration.ZERO);
        assertEquals(5, records.count());
        assertEquals(1, records.nextOffsets().size());
        assertEquals(new OffsetAndMetadata(5), records.nextOffsets().get(tp0));
        // Increment time asynchronously to clear timeouts in closing the consumer
        final ScheduledExecutorService exec = Executors.newSingleThreadScheduledExecutor();
        exec.scheduleAtFixedRate(() -> time.sleep(sessionTimeoutMs), 0L, 10L, TimeUnit.MILLISECONDS);
        consumer.close();
        exec.shutdownNow();
        exec.awaitTermination(5L, TimeUnit.SECONDS);
    }

    @ParameterizedTest
    @EnumSource(value = GroupProtocol.class, names = "CLASSIC")
    public void testPollThrowsInterruptExceptionIfInterrupted(GroupProtocol groupProtocol) {
        final ConsumerMetadata metadata = createMetadata(subscription);
        final MockClient client = new MockClient(time, metadata);

        initMetadata(client, Map.of(topic, 1));
        Node node = metadata.fetch().nodes().get(0);

        consumer = newConsumer(groupProtocol, time, client, subscription, metadata, assignor, false, groupInstanceId);
        consumer.subscribe(Set.of(topic), getConsumerRebalanceListener(consumer));
        prepareRebalance(client, node, assignor, List.of(tp0), null);

        consumer.updateAssignmentMetadataIfNeeded(time.timer(Long.MAX_VALUE));
        consumer.poll(Duration.ZERO);

        // interrupt the thread and call poll
        try {
            Thread.currentThread().interrupt();
            assertThrows(InterruptException.class, () -> consumer.poll(Duration.ZERO));
        } finally {
            // ensure this test doesn't exit with the thread still in an interrupted state
            // so that it cannot interfere with subsequent tests since this thread may be reused by JUnit
            Thread.interrupted();
        }
    }

    @ParameterizedTest
    @EnumSource(GroupProtocol.class)
    public void fetchResponseWithUnexpectedPartitionIsIgnored(GroupProtocol groupProtocol) {
        ConsumerMetadata metadata = createMetadata(subscription);
        MockClient client = new MockClient(time, metadata);

        initMetadata(client, Map.of(topic, 1));
        Node node = metadata.fetch().nodes().get(0);

        consumer = newConsumer(groupProtocol, time, client, subscription, metadata, assignor, true, groupInstanceId);
        consumer.subscribe(List.of(topic), getConsumerRebalanceListener(consumer));

        prepareRebalance(client, node, assignor, List.of(tp0), null);

        Map<TopicPartition, FetchInfo> fetches1 = new HashMap<>();
        fetches1.put(tp0, new FetchInfo(0, 1));
        fetches1.put(t2p0, new FetchInfo(0, 10)); // not assigned and not fetched
        client.prepareResponseFrom(fetchResponse(fetches1), node);

        consumer.updateAssignmentMetadataIfNeeded(time.timer(Long.MAX_VALUE));

        @SuppressWarnings("unchecked")
        ConsumerRecords<String, String> records = (ConsumerRecords<String, String>) consumer.poll(Duration.ZERO);
        assertEquals(0, records.count());
        assertEquals(0, records.nextOffsets().size());
    }

    /**
     * Verify that when a consumer changes its topic subscription its assigned partitions
     * do not immediately change, and the latest consumed offsets of its to-be-revoked
     * partitions are properly committed (when auto-commit is enabled).
     * Upon unsubscribing from subscribed topics the consumer subscription and assignment
     * are both updated right away but its consumed offsets are not auto committed.
     */
    // TODO: this test requires rebalance logic which is not yet implemented in the CONSUMER group protocol.
    //       Once it is implemented, this should use both group protocols.
    @ParameterizedTest
    @EnumSource(value = GroupProtocol.class, names = "CLASSIC")
    @SuppressWarnings("unchecked")
    public void testSubscriptionChangesWithAutoCommitEnabled(GroupProtocol groupProtocol) {
        ConsumerMetadata metadata = createMetadata(subscription);
        MockClient client = new MockClient(time, metadata);

        Map<String, Integer> tpCounts = new HashMap<>();
        tpCounts.put(topic, 1);
        tpCounts.put(topic2, 1);
        tpCounts.put(topic3, 1);
        initMetadata(client, tpCounts);
        Node node = metadata.fetch().nodes().get(0);

        ConsumerPartitionAssignor assignor = new RangeAssignor();

        consumer = newConsumer(groupProtocol, time, client, subscription, metadata, assignor, true, groupInstanceId);

        // initial subscription
        consumer.subscribe(Arrays.asList(topic, topic2), getConsumerRebalanceListener(consumer));

        // verify that subscription has changed but assignment is still unchanged
        assertEquals(2, consumer.subscription().size());
        assertTrue(consumer.subscription().contains(topic) && consumer.subscription().contains(topic2));
        assertTrue(consumer.assignment().isEmpty());

        // mock rebalance responses
        Node coordinator = prepareRebalance(client, node, assignor, Arrays.asList(tp0, t2p0), null);

        consumer.updateAssignmentMetadataIfNeeded(time.timer(Long.MAX_VALUE));
        consumer.poll(Duration.ZERO);

        // verify that subscription is still the same, and now assignment has caught up
        assertEquals(2, consumer.subscription().size());
        assertTrue(consumer.subscription().contains(topic) && consumer.subscription().contains(topic2));
        assertEquals(2, consumer.assignment().size());
        assertTrue(consumer.assignment().contains(tp0) && consumer.assignment().contains(t2p0));

        // mock a response to the outstanding fetch so that we have data available on the next poll
        Map<TopicPartition, FetchInfo> fetches1 = new HashMap<>();
        fetches1.put(tp0, new FetchInfo(0, 1));
        fetches1.put(t2p0, new FetchInfo(0, 10));
        client.respondFrom(fetchResponse(fetches1), node);
        client.poll(0, time.milliseconds());

        ConsumerRecords<String, String> records = (ConsumerRecords<String, String>) consumer.poll(Duration.ofMillis(1));

        // clear out the prefetch so it doesn't interfere with the rest of the test
        fetches1.put(tp0, new FetchInfo(1, 0));
        fetches1.put(t2p0, new FetchInfo(10, 0));
        client.respondFrom(fetchResponse(fetches1), node);
        client.poll(0, time.milliseconds());

        // verify that the fetch occurred as expected
        assertEquals(11, records.count());
        assertEquals(1L, consumer.position(tp0));
        assertEquals(10L, consumer.position(t2p0));

        // subscription change
        consumer.subscribe(Arrays.asList(topic, topic3), getConsumerRebalanceListener(consumer));

        // verify that subscription has changed but assignment is still unchanged
        assertEquals(2, consumer.subscription().size());
        assertTrue(consumer.subscription().contains(topic) && consumer.subscription().contains(topic3));
        assertEquals(2, consumer.assignment().size());
        assertTrue(consumer.assignment().contains(tp0) && consumer.assignment().contains(t2p0));

        // mock the offset commit response for to be revoked partitions
        Map<TopicPartition, Long> partitionOffsets1 = new HashMap<>();
        partitionOffsets1.put(tp0, 1L);
        partitionOffsets1.put(t2p0, 10L);
        AtomicBoolean commitReceived = prepareOffsetCommitResponse(client, coordinator, partitionOffsets1);

        // mock rebalance responses
        prepareRebalance(client, node, assignor, Arrays.asList(tp0, t3p0), coordinator);

        // mock a response to the next fetch from the new assignment
        Map<TopicPartition, FetchInfo> fetches2 = new HashMap<>();
        fetches2.put(tp0, new FetchInfo(1, 1));
        fetches2.put(t3p0, new FetchInfo(0, 100));
        client.prepareResponse(fetchResponse(fetches2));

        records = (ConsumerRecords<String, String>) consumer.poll(Duration.ofMillis(1));

        // verify that the fetch occurred as expected
        assertEquals(101, records.count());
        assertEquals(2L, consumer.position(tp0));
        assertEquals(100L, consumer.position(t3p0));
        assertEquals(2, records.nextOffsets().size());
        assertEquals(new OffsetAndMetadata(2), records.nextOffsets().get(tp0));
        assertEquals(new OffsetAndMetadata(100), records.nextOffsets().get(t3p0));

        // verify that the offset commits occurred as expected
        assertTrue(commitReceived.get());

        // verify that subscription is still the same, and now assignment has caught up
        assertEquals(2, consumer.subscription().size());
        assertTrue(consumer.subscription().contains(topic) && consumer.subscription().contains(topic3));
        assertEquals(2, consumer.assignment().size());
        assertTrue(consumer.assignment().contains(tp0) && consumer.assignment().contains(t3p0));

        consumer.unsubscribe();

        // verify that subscription and assignment are both cleared
        assertTrue(consumer.subscription().isEmpty());
        assertTrue(consumer.assignment().isEmpty());

        client.requests().clear();
    }

    /**
     * Verify that when a consumer changes its topic subscription its assigned partitions
     * do not immediately change, and the consumed offsets of its to-be-revoked partitions
     * are not committed (when auto-commit is disabled).
     * Upon unsubscribing from subscribed topics, the assigned partitions immediately
     * change but if auto-commit is disabled the consumer offsets are not committed.
     */
    // TODO: this test requires rebalance logic which is not yet implemented in the CONSUMER group protocol.
    //       Once it is implemented, this should use both group protocols.
    @ParameterizedTest
    @EnumSource(value = GroupProtocol.class, names = "CLASSIC")
    public void testSubscriptionChangesWithAutoCommitDisabled(GroupProtocol groupProtocol) {
        ConsumerMetadata metadata = createMetadata(subscription);
        MockClient client = new MockClient(time, metadata);

        Map<String, Integer> tpCounts = new HashMap<>();
        tpCounts.put(topic, 1);
        tpCounts.put(topic2, 1);
        initMetadata(client, tpCounts);
        Node node = metadata.fetch().nodes().get(0);

        ConsumerPartitionAssignor assignor = new RangeAssignor();

        consumer = newConsumer(groupProtocol, time, client, subscription, metadata, assignor, false, groupInstanceId);

        initializeSubscriptionWithSingleTopic(consumer, getConsumerRebalanceListener(consumer));

        // mock rebalance responses
        prepareRebalance(client, node, assignor, List.of(tp0), null);

        consumer.updateAssignmentMetadataIfNeeded(time.timer(Long.MAX_VALUE));
        consumer.poll(Duration.ZERO);

        // verify that subscription is still the same, and now assignment has caught up
        assertEquals(Set.of(topic), consumer.subscription());
        assertEquals(Set.of(tp0), consumer.assignment());

        consumer.poll(Duration.ZERO);

        // subscription change
        consumer.subscribe(Set.of(topic2), getConsumerRebalanceListener(consumer));

        // verify that subscription has changed but assignment is still unchanged
        assertEquals(Set.of(topic2), consumer.subscription());
        assertEquals(Set.of(tp0), consumer.assignment());

        // the auto commit is disabled, so no offset commit request should be sent
        for (ClientRequest req: client.requests())
            assertNotSame(ApiKeys.OFFSET_COMMIT, req.requestBuilder().apiKey());

        // subscription change
        consumer.unsubscribe();

        // verify that subscription and assignment are both updated
        assertEquals(Collections.emptySet(), consumer.subscription());
        assertEquals(Collections.emptySet(), consumer.assignment());

        // the auto commit is disabled, so no offset commit request should be sent
        for (ClientRequest req: client.requests())
            assertNotSame(ApiKeys.OFFSET_COMMIT, req.requestBuilder().apiKey());

        client.requests().clear();
    }

    // TODO: this test requires rebalance logic which is not yet implemented in the CONSUMER group protocol.
    //       Once it is implemented, this should use both group protocols.
    @ParameterizedTest
    @EnumSource(value = GroupProtocol.class, names = "CLASSIC")
    public void testUnsubscribeShouldTriggerPartitionsRevokedWithValidGeneration(GroupProtocol groupProtocol) {
        ConsumerMetadata metadata = createMetadata(subscription);
        MockClient client = new MockClient(time, metadata);

        initMetadata(client, Map.of(topic, 1));
        Node node = metadata.fetch().nodes().get(0);

        CooperativeStickyAssignor assignor = new CooperativeStickyAssignor();
        consumer = newConsumer(groupProtocol, time, client, subscription, metadata, assignor, false, groupInstanceId);

        initializeSubscriptionWithSingleTopic(consumer, getExceptionConsumerRebalanceListener());

        prepareRebalance(client, node, assignor, List.of(tp0), null);

        RuntimeException assignmentException = assertThrows(RuntimeException.class,
            () -> consumer.updateAssignmentMetadataIfNeeded(time.timer(Long.MAX_VALUE)));
        assertEquals(partitionAssigned + singleTopicPartition, assignmentException.getCause().getMessage());

        RuntimeException unsubscribeException = assertThrows(RuntimeException.class, consumer::unsubscribe);
        assertEquals(partitionRevoked + singleTopicPartition, unsubscribeException.getCause().getMessage());
    }

    // TODO: this test requires rebalance logic which is not yet implemented in the CONSUMER group protocol.
    //       Once it is implemented, this should use both group protocols.
    @ParameterizedTest
    @EnumSource(value = GroupProtocol.class, names = "CLASSIC")
    public void testUnsubscribeShouldTriggerPartitionsLostWithNoGeneration(GroupProtocol groupProtocol) throws Exception {
        ConsumerMetadata metadata = createMetadata(subscription);
        MockClient client = new MockClient(time, metadata);

        initMetadata(client, Map.of(topic, 1));
        Node node = metadata.fetch().nodes().get(0);

        CooperativeStickyAssignor assignor = new CooperativeStickyAssignor();
        consumer = newConsumer(groupProtocol, time, client, subscription, metadata, assignor, false, groupInstanceId);

        initializeSubscriptionWithSingleTopic(consumer, getExceptionConsumerRebalanceListener());
        Node coordinator = prepareRebalance(client, node, assignor, List.of(tp0), null);

        RuntimeException assignException = assertThrows(RuntimeException.class,
            () -> consumer.updateAssignmentMetadataIfNeeded(time.timer(Long.MAX_VALUE)));
        assertEquals(partitionAssigned + singleTopicPartition, assignException.getCause().getMessage());

        AtomicBoolean heartbeatReceived = prepareHeartbeatResponse(client, coordinator, Errors.UNKNOWN_MEMBER_ID);

        time.sleep(heartbeatIntervalMs);
        TestUtils.waitForCondition(heartbeatReceived::get, "Heartbeat response did not occur within timeout.");

        RuntimeException unsubscribeException = assertThrows(RuntimeException.class, consumer::unsubscribe);
        assertEquals(partitionLost + singleTopicPartition, unsubscribeException.getCause().getMessage());
    }

    private void initializeSubscriptionWithSingleTopic(KafkaConsumer<?, ?> consumer,
                                                       ConsumerRebalanceListener consumerRebalanceListener) {
        consumer.subscribe(Set.of(topic), consumerRebalanceListener);
        // verify that subscription has changed but assignment is still unchanged
        assertEquals(Set.of(topic), consumer.subscription());
        assertEquals(Collections.emptySet(), consumer.assignment());
    }

    @ParameterizedTest
    @EnumSource(GroupProtocol.class)
    @SuppressWarnings("unchecked")
    public void testManualAssignmentChangeWithAutoCommitEnabled(GroupProtocol groupProtocol) {
        ConsumerMetadata metadata = createMetadata(subscription);
        MockClient client = new MockClient(time, metadata);

        Map<String, Integer> tpCounts = new HashMap<>();
        tpCounts.put(topic, 1);
        tpCounts.put(topic2, 1);
        initMetadata(client, tpCounts);
        Node node = metadata.fetch().nodes().get(0);

        ConsumerPartitionAssignor assignor = new RangeAssignor();

        client.prepareResponseFrom(FindCoordinatorResponse.prepareResponse(Errors.NONE, groupId, node), node);
        consumer = newConsumer(groupProtocol, time, client, subscription, metadata, assignor, true, groupInstanceId);

        // lookup coordinator
        Node coordinator = new Node(Integer.MAX_VALUE - node.id(), node.host(), node.port());

        // manual assignment
        consumer.assign(Set.of(tp0));
        consumer.seekToBeginning(Set.of(tp0));

        // fetch offset for one topic
        client.prepareResponseFrom(offsetResponse(Map.of(tp0, 0L), Errors.NONE), coordinator);
        assertEquals(0, consumer.committed(Set.of(tp0)).get(tp0).offset());

        // verify that assignment immediately changes
        assertEquals(Set.of(tp0), consumer.assignment());

        // there shouldn't be any need to look up the coordinator or fetch committed offsets.
        // we just look up the starting position and send the record fetch.
        client.prepareResponse(listOffsetsResponse(Map.of(tp0, 10L)));
        client.prepareResponse(fetchResponse(tp0, 10L, 1));

        ConsumerRecords<String, String> records = (ConsumerRecords<String, String>) consumer.poll(Duration.ofMillis(100));

        assertEquals(1, records.count());
        assertEquals(11L, consumer.position(tp0));
        assertEquals(1, records.nextOffsets().size());
        assertEquals(new OffsetAndMetadata(11L), records.nextOffsets().get(tp0));

        // mock the offset commit response for to be revoked partitions
        AtomicBoolean commitReceived = prepareOffsetCommitResponse(client, coordinator, tp0, 11);

        // new manual assignment
        consumer.assign(Set.of(t2p0));

        // verify that assignment immediately changes
        assertEquals(Set.of(t2p0), consumer.assignment());
        // verify that the offset commits occurred as expected
        assertTrue(commitReceived.get());

        client.requests().clear();
    }

    @ParameterizedTest
    @EnumSource(GroupProtocol.class)
    public void testManualAssignmentChangeWithAutoCommitDisabled(GroupProtocol groupProtocol) {
        ConsumerMetadata metadata = createMetadata(subscription);
        MockClient client = new MockClient(time, metadata);

        Map<String, Integer> tpCounts = new HashMap<>();
        tpCounts.put(topic, 1);
        tpCounts.put(topic2, 1);
        initMetadata(client, tpCounts);
        Node node = metadata.fetch().nodes().get(0);

        ConsumerPartitionAssignor assignor = new RangeAssignor();

        client.prepareResponseFrom(FindCoordinatorResponse.prepareResponse(Errors.NONE, groupId, node), node);
        consumer = newConsumer(groupProtocol, time, client, subscription, metadata, assignor, false, groupInstanceId);

        // lookup coordinator
        Node coordinator = new Node(Integer.MAX_VALUE - node.id(), node.host(), node.port());

        // manual assignment
        consumer.assign(Set.of(tp0));
        consumer.seekToBeginning(Set.of(tp0));

        // fetch offset for one topic
        client.prepareResponseFrom(
                offsetResponse(Map.of(tp0, 0L), Errors.NONE),
                coordinator);
        assertEquals(0, consumer.committed(Set.of(tp0)).get(tp0).offset());

        // verify that assignment immediately changes
        assertEquals(Set.of(tp0), consumer.assignment());

        // there shouldn't be any need to look up the coordinator or fetch committed offsets.
        // we just look up the starting position and send the record fetch.
        client.prepareResponse(listOffsetsResponse(Map.of(tp0, 10L)));
        client.prepareResponse(fetchResponse(tp0, 10L, 1));

        @SuppressWarnings("unchecked")
        ConsumerRecords<String, String> records = (ConsumerRecords<String, String>) consumer.poll(Duration.ofMillis(1));
        assertEquals(1, records.count());
        assertEquals(11L, consumer.position(tp0));
        assertEquals(1, records.nextOffsets().size());
        assertEquals(new OffsetAndMetadata(11L), records.nextOffsets().get(tp0));

        // new manual assignment
        consumer.assign(Set.of(t2p0));

        // verify that assignment immediately changes
        assertEquals(Set.of(t2p0), consumer.assignment());

        // the auto commit is disabled, so no offset commit request should be sent
        for (ClientRequest req : client.requests())
            assertNotSame(ApiKeys.OFFSET_COMMIT, req.requestBuilder().apiKey());

        client.requests().clear();
    }

    @ParameterizedTest
    @EnumSource(GroupProtocol.class)
    public void testOffsetOfPausedPartitions(GroupProtocol groupProtocol) {
        ConsumerMetadata metadata = createMetadata(subscription);
        MockClient client = new MockClient(time, metadata);

        initMetadata(client, Map.of(topic, 2));
        Node node = metadata.fetch().nodes().get(0);

        ConsumerPartitionAssignor assignor = new RangeAssignor();

        client.prepareResponseFrom(FindCoordinatorResponse.prepareResponse(Errors.NONE, groupId, node), node);
        consumer = newConsumer(groupProtocol, time, client, subscription, metadata, assignor, true, groupInstanceId);

        // lookup coordinator
        Node coordinator = new Node(Integer.MAX_VALUE - node.id(), node.host(), node.port());

        // manual assignment
        Set<TopicPartition> partitions = Set.of(tp0, tp1);
        consumer.assign(partitions);
        // verify consumer's assignment
        assertEquals(partitions, consumer.assignment());

        consumer.pause(partitions);
        consumer.seekToEnd(partitions);

        // fetch and verify committed offset of two partitions
        Map<TopicPartition, Long> offsets = new HashMap<>();
        offsets.put(tp0, 0L);
        offsets.put(tp1, 0L);

        client.prepareResponseFrom(offsetResponse(offsets, Errors.NONE), coordinator);
        assertEquals(0, consumer.committed(Set.of(tp0)).get(tp0).offset());

        offsets.remove(tp0);
        offsets.put(tp1, 0L);
        client.prepareResponseFrom(offsetResponse(offsets, Errors.NONE), coordinator);
        assertEquals(0, consumer.committed(Set.of(tp1)).get(tp1).offset());

        // fetch and verify consumer's position in the two partitions
        final Map<TopicPartition, Long> offsetResponse = new HashMap<>();
        offsetResponse.put(tp0, 3L);
        offsetResponse.put(tp1, 3L);
        client.prepareResponse(listOffsetsResponse(offsetResponse));
        assertEquals(3L, consumer.position(tp0));
        assertEquals(3L, consumer.position(tp1));

        client.requests().clear();
        consumer.unsubscribe();
    }

    @ParameterizedTest
    @EnumSource(GroupProtocol.class)
    public void testPollWithNoSubscription(GroupProtocol groupProtocol) {
        consumer = newConsumer(groupProtocol, null);
        assertThrows(IllegalStateException.class, () -> consumer.poll(Duration.ZERO));
    }

    @ParameterizedTest
    @EnumSource(GroupProtocol.class)
    public void testPollWithEmptySubscription(GroupProtocol groupProtocol) {
        consumer = newConsumer(groupProtocol, groupId);
        consumer.subscribe(Collections.emptyList());
        assertThrows(IllegalStateException.class, () -> consumer.poll(Duration.ZERO));
    }

    @ParameterizedTest
    @EnumSource(GroupProtocol.class)
    public void testPollWithEmptyUserAssignment(GroupProtocol groupProtocol) {
        consumer = newConsumer(groupProtocol, groupId);
        consumer.assign(Collections.emptySet());
        assertThrows(IllegalStateException.class, () -> consumer.poll(Duration.ZERO));
    }

    // TODO: this test references RPCs to be sent that are not part of the CONSUMER group protocol.
    //       We are deferring any attempts at generalizing this test for both group protocols to the future.
    @ParameterizedTest
    @EnumSource(value = GroupProtocol.class, names = "CLASSIC")
    public void testGracefulClose(GroupProtocol groupProtocol) throws Exception {
        Map<TopicPartition, Errors> response = new HashMap<>();
        response.put(tp0, Errors.NONE);
        OffsetCommitResponse commitResponse = offsetCommitResponse(response);
        LeaveGroupResponse leaveGroupResponse = new LeaveGroupResponse(new LeaveGroupResponseData().setErrorCode(Errors.NONE.code()));
        FetchResponse closeResponse = FetchResponse.of(Errors.NONE, 0, INVALID_SESSION_ID, new LinkedHashMap<>(), List.of());
        consumerCloseTest(groupProtocol, 5000, Arrays.asList(commitResponse, leaveGroupResponse, closeResponse), 0, false);
    }

    // TODO: this test references RPCs to be sent that are not part of the CONSUMER group protocol.
    //       We are deferring any attempts at generalizing this test for both group protocols to the future.
    @ParameterizedTest
    @EnumSource(value = GroupProtocol.class, names = "CLASSIC")
    public void testCloseTimeoutDueToNoResponseForCloseFetchRequest(GroupProtocol groupProtocol) throws Exception {
        Map<TopicPartition, Errors> response = new HashMap<>();
        response.put(tp0, Errors.NONE);
        OffsetCommitResponse commitResponse = offsetCommitResponse(response);
        LeaveGroupResponse leaveGroupResponse = new LeaveGroupResponse(new LeaveGroupResponseData().setErrorCode(Errors.NONE.code()));
        final List<AbstractResponse> serverResponsesWithoutCloseResponse = Arrays.asList(commitResponse, leaveGroupResponse);

        // to ensure timeout due to no response for fetcher close request, we will ensure that we have successful
        // response from server for first two requests and the test is configured to wait for duration which is greater
        // than configured timeout.
        final int closeTimeoutMs = 5000;
        final int waitForCloseCompletionMs = closeTimeoutMs + 1000;
        consumerCloseTest(groupProtocol, closeTimeoutMs, serverResponsesWithoutCloseResponse, waitForCloseCompletionMs, false);
    }

    // TODO: this test requires rebalance logic which is not yet implemented in the CONSUMER group protocol.
    //       Once it is implemented, this should use both group protocols.
    @ParameterizedTest
    @EnumSource(value = GroupProtocol.class, names = "CLASSIC")
    public void testCloseTimeout(GroupProtocol groupProtocol) throws Exception {
        consumerCloseTest(groupProtocol, 5000, Collections.emptyList(), 5000, false);
    }

    // TODO: this test requires rebalance logic which is not yet implemented in the CONSUMER group protocol.
    //       Once it is implemented, this should use both group protocols.
    @ParameterizedTest
    @EnumSource(value = GroupProtocol.class, names = "CLASSIC")
    public void testLeaveGroupTimeout(GroupProtocol groupProtocol) throws Exception {
        Map<TopicPartition, Errors> response = new HashMap<>();
        response.put(tp0, Errors.NONE);
        OffsetCommitResponse commitResponse = offsetCommitResponse(response);
        consumerCloseTest(groupProtocol, 5000, List.of(commitResponse), 5000, false);
    }

    // TODO: this test requires rebalance logic which is not yet implemented in the CONSUMER group protocol.
    //       Once it is implemented, this should use both group protocols.
    @ParameterizedTest
    @EnumSource(value = GroupProtocol.class, names = "CLASSIC")
    public void testCloseNoWait(GroupProtocol groupProtocol) throws Exception {
        consumerCloseTest(groupProtocol, 0, Collections.emptyList(), 0, false);
    }

    // TODO: this test requires rebalance logic which is not yet implemented in the CONSUMER group protocol.
    //       Once it is implemented, this should use both group protocols.
    @ParameterizedTest
    @EnumSource(value = GroupProtocol.class, names = "CLASSIC")
    public void testCloseInterrupt(GroupProtocol groupProtocol) throws Exception {
        consumerCloseTest(groupProtocol, Long.MAX_VALUE, Collections.emptyList(), 0, true);
    }

    @ParameterizedTest
    @EnumSource(GroupProtocol.class)
    public void testCloseShouldBeIdempotent(GroupProtocol groupProtocol) {
        ConsumerMetadata metadata = createMetadata(subscription);
        MockClient client = spy(new MockClient(time, metadata));
        initMetadata(client, Map.of(topic, 1));

        consumer = newConsumer(groupProtocol, time, client, subscription, metadata, assignor, false, groupInstanceId);

        consumer.close(CloseOptions.timeout(Duration.ZERO));
        consumer.close(CloseOptions.timeout(Duration.ZERO));

        // verify that the call is idempotent by checking that the network client is only closed once.
        verify(client).close();
    }

    @SuppressWarnings("resource")
    @ParameterizedTest
    @EnumSource(GroupProtocol.class)
    public void testOperationsBySubscribingConsumerWithDefaultGroupId(GroupProtocol groupProtocol) {
        assertThrows(InvalidConfigurationException.class,
            () -> newConsumer(groupProtocol, null, Optional.of(true)),
            "Expected an InvalidConfigurationException");

        try (KafkaConsumer<byte[], byte[]> consumer = newConsumer(groupProtocol, null)) {
            assertThrows(InvalidGroupIdException.class, () -> consumer.subscribe(Set.of(topic)));
        }

        try (KafkaConsumer<byte[], byte[]> consumer = newConsumer(groupProtocol, null)) {
            assertThrows(InvalidGroupIdException.class, () -> consumer.committed(Set.of(tp0)).get(tp0));
        }

        try (KafkaConsumer<byte[], byte[]> consumer = newConsumer(groupProtocol, null)) {
            assertThrows(InvalidGroupIdException.class, consumer::commitAsync);
        }

        try (KafkaConsumer<byte[], byte[]> consumer = newConsumer(groupProtocol, null)) {
            assertThrows(InvalidGroupIdException.class, consumer::commitSync);
        }
    }

    @ParameterizedTest
    @EnumSource(GroupProtocol.class)
    public void testOperationsByAssigningConsumerWithDefaultGroupId(GroupProtocol groupProtocol) {
        try (KafkaConsumer<byte[], byte[]> consumer = newConsumer(groupProtocol, null)) {
            consumer.assign(Set.of(tp0));

            assertThrows(InvalidGroupIdException.class, () -> consumer.committed(Set.of(tp0)).get(tp0));
            assertThrows(InvalidGroupIdException.class, consumer::commitAsync);
            assertThrows(InvalidGroupIdException.class, consumer::commitSync);
        }
    }

    @ParameterizedTest
    @EnumSource(GroupProtocol.class)
    public void testMetricConfigRecordingLevelInfo(GroupProtocol groupProtocol) {
        Properties props = new Properties();
        props.put(ConsumerConfig.GROUP_PROTOCOL_CONFIG, groupProtocol.name());
        props.put(ConsumerConfig.BOOTSTRAP_SERVERS_CONFIG, "localhost:9000");
        KafkaConsumer<byte[], byte[]> consumer = newConsumer(props, new ByteArrayDeserializer(), new ByteArrayDeserializer());
        assertEquals(Sensor.RecordingLevel.INFO, consumer.metricsRegistry().config().recordLevel());
        consumer.close(CloseOptions.timeout(Duration.ZERO));

        props.put(ConsumerConfig.METRICS_RECORDING_LEVEL_CONFIG, "DEBUG");
        KafkaConsumer<byte[], byte[]> consumer2 = newConsumer(props, new ByteArrayDeserializer(), new ByteArrayDeserializer());
        assertEquals(Sensor.RecordingLevel.DEBUG, consumer2.metricsRegistry().config().recordLevel());
        consumer2.close(CloseOptions.timeout(Duration.ZERO));
    }

    // TODO: this test references RPCs to be sent that are not part of the CONSUMER group protocol.
    //       We are deferring any attempts at generalizing this test for both group protocols to the future.
    @ParameterizedTest
    @EnumSource(value = GroupProtocol.class, names = "CLASSIC")
    @SuppressWarnings("unchecked")
    public void testShouldAttemptToRejoinGroupAfterSyncGroupFailed(GroupProtocol groupProtocol) throws Exception {
        ConsumerMetadata metadata = createMetadata(subscription);
        MockClient client = new MockClient(time, metadata);

        initMetadata(client, Map.of(topic, 1));
        Node node = metadata.fetch().nodes().get(0);

        consumer = newConsumer(groupProtocol, time, client, subscription, metadata, assignor, false, groupInstanceId);
        consumer.subscribe(Set.of(topic), getConsumerRebalanceListener(consumer));
        client.prepareResponseFrom(FindCoordinatorResponse.prepareResponse(Errors.NONE, groupId, node), node);
        Node coordinator = new Node(Integer.MAX_VALUE - node.id(), node.host(), node.port());


        client.prepareResponseFrom(joinGroupFollowerResponse(assignor, 1, memberId, leaderId, Errors.NONE), coordinator);
        client.prepareResponseFrom(syncGroupResponse(List.of(tp0), Errors.NONE), coordinator);

        client.prepareResponseFrom(fetchResponse(tp0, 0, 1), node);
        client.prepareResponseFrom(fetchResponse(tp0, 1, 0), node);

        consumer.updateAssignmentMetadataIfNeeded(time.timer(Long.MAX_VALUE));
        consumer.poll(Duration.ZERO);

        // heartbeat fails due to rebalance in progress
        client.prepareResponseFrom(body -> true, new HeartbeatResponse(
            new HeartbeatResponseData().setErrorCode(Errors.REBALANCE_IN_PROGRESS.code())), coordinator);

        // join group
        final ByteBuffer byteBuffer = ConsumerProtocol.serializeSubscription(new ConsumerPartitionAssignor.Subscription(List.of(topic)));

        // This member becomes the leader
        final JoinGroupResponse leaderResponse = new JoinGroupResponse(
                new JoinGroupResponseData()
                        .setErrorCode(Errors.NONE.code())
                        .setGenerationId(1).setProtocolName(assignor.name())
                        .setLeader(memberId).setMemberId(memberId)
                        .setMembers(List.of(
                                new JoinGroupResponseData.JoinGroupResponseMember()
                                        .setMemberId(memberId)
                                        .setMetadata(byteBuffer.array())
                                )
                        ),
                ApiKeys.JOIN_GROUP.latestVersion()
        );

        client.prepareResponseFrom(leaderResponse, coordinator);

        // sync group fails due to disconnect
        client.prepareResponseFrom(syncGroupResponse(List.of(tp0), Errors.NONE), coordinator, true);

        // should try and find the new coordinator
        client.prepareResponseFrom(FindCoordinatorResponse.prepareResponse(Errors.NONE, groupId, node), node);

        // rejoin group
        client.prepareResponseFrom(joinGroupFollowerResponse(assignor, 1, memberId, leaderId, Errors.NONE), coordinator);
        client.prepareResponseFrom(syncGroupResponse(List.of(tp0), Errors.NONE), coordinator);

        client.prepareResponseFrom(body -> body instanceof FetchRequest
            && ((FetchRequest) body).fetchData(topicNames).containsKey(new TopicIdPartition(topicId, tp0)), fetchResponse(tp0, 1, 1), node);
        time.sleep(heartbeatIntervalMs);
        Thread.sleep(heartbeatIntervalMs);
        consumer.updateAssignmentMetadataIfNeeded(time.timer(Long.MAX_VALUE));
        final ConsumerRecords<String, String> records = (ConsumerRecords<String, String>) consumer.poll(Duration.ZERO);
        assertFalse(records.isEmpty());
        assertFalse(records.nextOffsets().isEmpty());
    }

    private void consumerCloseTest(GroupProtocol groupProtocol,
                                   final long closeTimeoutMs,
                                   List<? extends AbstractResponse> responses,
                                   long waitMs,
                                   boolean interrupt) throws Exception {
        ConsumerMetadata metadata = createMetadata(subscription);
        MockClient client = new MockClient(time, metadata);

        initMetadata(client, Map.of(topic, 1));
        Node node = metadata.fetch().nodes().get(0);

        final KafkaConsumer<String, String> consumer = newConsumer(groupProtocol, time, client, subscription, metadata, assignor, false, Optional.empty());
        consumer.subscribe(Set.of(topic), getConsumerRebalanceListener(consumer));
        Node coordinator = prepareRebalance(client, node, assignor, List.of(tp0), null);

        client.prepareMetadataUpdate(RequestTestUtils.metadataUpdateWithIds(1, Map.of(topic, 1), topicIds));

        consumer.updateAssignmentMetadataIfNeeded(time.timer(Long.MAX_VALUE));

        // Poll with responses
        client.prepareResponseFrom(fetchResponse(tp0, 0, 1), node);
        client.prepareResponseFrom(fetchResponse(tp0, 1, 0), node);
        consumer.poll(Duration.ZERO);

        // Initiate close() after a commit request on another thread.
        // Kafka consumer is single-threaded, but the implementation allows calls on a
        // different thread as long as the calls are not executed concurrently. So this is safe.
        ExecutorService executor = Executors.newSingleThreadExecutor();
        final AtomicReference<Exception> closeException = new AtomicReference<>();
        try {
            Future<?> future = executor.submit(() -> {
                consumer.commitAsync();
                try {
                    consumer.close(CloseOptions.timeout(Duration.ofMillis(closeTimeoutMs)));
                } catch (Exception e) {
                    closeException.set(e);
                }
            });

            // Close task should not complete until commit succeeds or close times out
            // if close timeout is not zero.
            try {
                future.get(100, TimeUnit.MILLISECONDS);
                if (closeTimeoutMs != 0)
                    fail("Close completed without waiting for commit or leave response");
            } catch (TimeoutException swallow) {
                // Expected exception
            }

            // Ensure close has started and queued at least one more request after commitAsync.
            //
            // Close enqueues two requests, but second is enqueued only after first has succeeded. First is
            // LEAVE_GROUP as part of coordinator close and second is FETCH with epoch=FINAL_EPOCH. At this stage
            // we expect only the first one to have been requested. Hence, waiting for total 2 requests, one for
            // commit and another for LEAVE_GROUP.
            client.waitForRequests(2, 1000);

            // In graceful mode, commit response results in close() completing immediately without a timeout
            // In non-graceful mode, close() times out without an exception even though commit response is pending
            int nonCloseRequests = 1;
            for (int i = 0; i < responses.size(); i++) {
                client.waitForRequests(1, 1000);
                if (i == responses.size() - 1 && responses.get(i) instanceof FetchResponse) {
                    // last request is the close session request which is sent to the leader of the partition.
                    client.respondFrom(responses.get(i), node);
                } else {
                    client.respondFrom(responses.get(i), coordinator);
                }
                if (i < nonCloseRequests) {
                    // the close request should not complete until non-close requests (commit requests) have completed.
                    assertThrows(TimeoutException.class, () -> future.get(100, TimeUnit.MILLISECONDS));
                }
            }

            if (waitMs > 0)
                time.sleep(waitMs);
            if (interrupt) {
                assertTrue(future.cancel(true), "Close terminated prematurely");

                TestUtils.waitForCondition(
                    () -> closeException.get() != null, "InterruptException did not occur within timeout.");

                assertInstanceOf(InterruptException.class, closeException.get(), "Expected exception not thrown " + closeException);
            } else {
                future.get(closeTimeoutMs, TimeUnit.MILLISECONDS); // Should succeed without TimeoutException or ExecutionException
                assertNull(closeException.get(), "Unexpected exception during close");
            }
        } finally {
            executor.shutdownNow();
        }
    }

    // TODO: this test requires topic metadata logic which is not yet implemented in the CONSUMER group protocol.
    //       Once it is implemented, this should use both group protocols.
    @ParameterizedTest
    @EnumSource(value = GroupProtocol.class, names = "CLASSIC")
    public void testPartitionsForNonExistingTopic(GroupProtocol groupProtocol) {
        ConsumerMetadata metadata = createMetadata(subscription);
        MockClient client = new MockClient(time, metadata);

        initMetadata(client, Map.of(topic, 1));
        Cluster cluster = metadata.fetch();

        MetadataResponse updateResponse = RequestTestUtils.metadataResponse(cluster.nodes(),
            cluster.clusterResource().clusterId(),
            cluster.controller().id(),
            Collections.emptyList());
        client.prepareResponse(updateResponse);

        KafkaConsumer<String, String> consumer = newConsumer(groupProtocol, time, client, subscription, metadata, assignor, true, groupInstanceId);
        assertEquals(Collections.emptyList(), consumer.partitionsFor("non-exist-topic"));
    }

    // TODO: this test requires topic metadata logic which is not yet implemented in the CONSUMER group protocol.
    //       Once it is implemented, this should use both group protocols.
    @ParameterizedTest
    @EnumSource(value = GroupProtocol.class, names = "CLASSIC")
    public void testPartitionsForAuthenticationFailure(GroupProtocol groupProtocol) {
        try (final KafkaConsumer<String, String> consumer = consumerWithPendingAuthenticationError(groupProtocol)) {
            assertThrows(AuthenticationException.class, () -> consumer.partitionsFor("some other topic"));
        }
    }

    // TODO: this test triggers a bug with the CONSUMER group protocol implementation.
    //       The bug will be investigated and fixed so this test can use both group protocols.
    @ParameterizedTest
    @EnumSource(value = GroupProtocol.class, names = "CLASSIC")
    public void testBeginningOffsetsAuthenticationFailure(GroupProtocol groupProtocol) {
        try (final KafkaConsumer<String, String> consumer = consumerWithPendingAuthenticationError(groupProtocol)) {
            assertThrows(AuthenticationException.class, () -> consumer.beginningOffsets(Set.of(tp0)));
        }
    }

    // TODO: this test triggers a bug with the CONSUMER group protocol implementation.
    //       The bug will be investigated and fixed so this test can use both group protocols.
    @ParameterizedTest
    @EnumSource(value = GroupProtocol.class, names = "CLASSIC")
    public void testEndOffsetsAuthenticationFailure(GroupProtocol groupProtocol) {
        try (final KafkaConsumer<String, String> consumer = consumerWithPendingAuthenticationError(groupProtocol)) {
            assertThrows(AuthenticationException.class, () -> consumer.endOffsets(Set.of(tp0)));
        }
    }

    @ParameterizedTest
    @EnumSource(GroupProtocol.class)
    public void testPollAuthenticationFailure(GroupProtocol groupProtocol) throws InterruptedException {
        final KafkaConsumer<String, String> consumer = consumerWithPendingAuthenticationError(groupProtocol);
        consumer.subscribe(Set.of(topic));

        if (groupProtocol == GroupProtocol.CONSUMER) {
            // New consumer poll(ZERO) needs to wait for the event added by a call to poll, to be processed
            // by the background thread, so it can realize there is authentication fail  and then
            // throw the AuthenticationException
            assertPollEventuallyThrows(consumer, AuthenticationException.class,
                    "this consumer was not able to discover metadata errors during continuous polling.");
        } else {
            assertThrows(AuthenticationException.class, () -> consumer.poll(Duration.ZERO));
        }
    }

    // TODO: this test triggers a bug with the CONSUMER group protocol implementation.
    //       The bug will be investigated and fixed so this test can use both group protocols.
    @ParameterizedTest
    @EnumSource(value = GroupProtocol.class, names = "CLASSIC")
    public void testOffsetsForTimesAuthenticationFailure(GroupProtocol groupProtocol) {
        try (final KafkaConsumer<String, String> consumer = consumerWithPendingAuthenticationError(groupProtocol)) {
            assertThrows(AuthenticationException.class, () -> consumer.offsetsForTimes(Map.of(tp0, 0L)));
        }
    }

    // TODO: this test triggers a bug with the CONSUMER group protocol implementation.
    //       The bug will be investigated and fixed so this test can use both group protocols.
    @ParameterizedTest
    @EnumSource(value = GroupProtocol.class, names = "CLASSIC")
    public void testCommitSyncAuthenticationFailure(GroupProtocol groupProtocol) {
        try (final KafkaConsumer<String, String> consumer = consumerWithPendingAuthenticationError(groupProtocol)) {
            Map<TopicPartition, OffsetAndMetadata> offsets = Map.of(tp0, new OffsetAndMetadata(10L));
            assertThrows(AuthenticationException.class, () -> consumer.commitSync(offsets));
        }
    }

    // TODO: this test triggers a bug with the CONSUMER group protocol implementation.
    //       The bug will be investigated and fixed so this test can use both group protocols.
    @ParameterizedTest
    @EnumSource(value = GroupProtocol.class, names = "CLASSIC")
    public void testCommittedAuthenticationFailure(GroupProtocol groupProtocol) {
        try (final KafkaConsumer<String, String> consumer = consumerWithPendingAuthenticationError(groupProtocol)) {
            assertThrows(AuthenticationException.class, () -> consumer.committed(Set.of(tp0)).get(tp0));
        }
    }

    @ParameterizedTest
    @EnumSource(value = GroupProtocol.class)
    public void testMeasureCommitSyncDurationOnFailure(GroupProtocol groupProtocol) {
        try (final KafkaConsumer<String, String> consumer
            = consumerWithPendingError(groupProtocol, new MockTime(Duration.ofSeconds(1).toMillis()))) {

            try {
                consumer.commitSync(Map.of(tp0, new OffsetAndMetadata(10L)));
            } catch (final RuntimeException swallow) {
                // swallow
            }

            final Metric metric = consumer.metrics()
                    .get(consumer.metricsRegistry().metricName("commit-sync-time-ns-total", "consumer-metrics"));
            assertTrue((Double) metric.metricValue() >= Duration.ofMillis(999).toNanos());
        }
    }

    @ParameterizedTest
    @EnumSource(value = GroupProtocol.class)
    public void testMeasureCommitSyncDuration(GroupProtocol groupProtocol) {
        Time time = new MockTime(Duration.ofSeconds(1).toMillis());
        SubscriptionState subscription = new SubscriptionState(new LogContext(),
            AutoOffsetResetStrategy.EARLIEST);
        ConsumerMetadata metadata = createMetadata(subscription);
        MockClient client = new MockClient(time, metadata);
        initMetadata(client, Map.of(topic, 2));
        Node node = metadata.fetch().nodes().get(0);
        KafkaConsumer<String, String> consumer = newConsumer(groupProtocol, time, client, subscription, metadata,
            assignor, true, groupInstanceId);
        consumer.assign(List.of(tp0));

        client.prepareResponseFrom(
            FindCoordinatorResponse.prepareResponse(Errors.NONE, groupId, node), node);
        Node coordinator = new Node(Integer.MAX_VALUE - node.id(), node.host(), node.port());
        client.prepareResponseFrom(
            offsetCommitResponse(Map.of(tp0, Errors.NONE)),
            coordinator
        );

        consumer.commitSync(Map.of(tp0, new OffsetAndMetadata(10L)));

        final Metric metric = consumer.metrics()
            .get(consumer.metricsRegistry().metricName("commit-sync-time-ns-total", "consumer-metrics"));
        assertTrue((Double) metric.metricValue() >= Duration.ofMillis(999).toNanos());
    }

    @ParameterizedTest
    @EnumSource(value = GroupProtocol.class)
    public void testMeasureCommittedDurationOnFailure(GroupProtocol groupProtocol) {
        try (final KafkaConsumer<String, String> consumer
            = consumerWithPendingError(groupProtocol, new MockTime(Duration.ofSeconds(1).toMillis()))) {

            try {
                consumer.committed(Set.of(tp0));
            } catch (final RuntimeException swallow) {
                // swallow
            }

            final Metric metric = consumer.metrics()
                    .get(consumer.metricsRegistry().metricName("committed-time-ns-total", "consumer-metrics"));
            assertTrue((Double) metric.metricValue() >= Duration.ofMillis(999).toNanos());
        }
    }

    @ParameterizedTest
    @EnumSource(value = GroupProtocol.class)
    public void testMeasureCommittedDuration(GroupProtocol groupProtocol) {
        long offset1 = 10000;
        Time time = new MockTime(Duration.ofSeconds(1).toMillis());
        SubscriptionState subscription = new SubscriptionState(new LogContext(),
            AutoOffsetResetStrategy.EARLIEST);
        ConsumerMetadata metadata = createMetadata(subscription);
        MockClient client = new MockClient(time, metadata);
        initMetadata(client, Map.of(topic, 2));
        Node node = metadata.fetch().nodes().get(0);
        KafkaConsumer<String, String> consumer = newConsumer(groupProtocol, time, client, subscription, metadata,
            assignor, true, groupInstanceId);
        consumer.assign(List.of(tp0));

        // lookup coordinator
        client.prepareResponseFrom(
            FindCoordinatorResponse.prepareResponse(Errors.NONE, groupId, node), node);
        Node coordinator = new Node(Integer.MAX_VALUE - node.id(), node.host(), node.port());

        // fetch offset for one topic
        client.prepareResponseFrom(
            offsetResponse(Map.of(tp0, offset1), Errors.NONE), coordinator);

        consumer.committed(Set.of(tp0)).get(tp0).offset();

        final Metric metric = consumer.metrics()
            .get(consumer.metricsRegistry().metricName("committed-time-ns-total", "consumer-metrics"));
        assertTrue((Double) metric.metricValue() >= Duration.ofMillis(999).toNanos());
    }

    // TODO: this test requires rebalance logic which is not yet implemented in the CONSUMER group protocol.
    //       Once it is implemented, this should use both group protocols.
    @ParameterizedTest
    @EnumSource(value = GroupProtocol.class, names = "CLASSIC")
    public void testRebalanceException(GroupProtocol groupProtocol) {
        ConsumerMetadata metadata = createMetadata(subscription);
        MockClient client = new MockClient(time, metadata);

        initMetadata(client, Map.of(topic, 1));
        Node node = metadata.fetch().nodes().get(0);

        KafkaConsumer<String, String> consumer = newConsumer(groupProtocol, time, client, subscription, metadata, assignor, true, groupInstanceId);

        consumer.subscribe(Set.of(topic), getExceptionConsumerRebalanceListener());
        Node coordinator = new Node(Integer.MAX_VALUE - node.id(), node.host(), node.port());

        client.prepareResponseFrom(FindCoordinatorResponse.prepareResponse(Errors.NONE, groupId, node), node);
        client.prepareResponseFrom(joinGroupFollowerResponse(assignor, 1, memberId, leaderId, Errors.NONE), coordinator);
        client.prepareResponseFrom(syncGroupResponse(List.of(tp0), Errors.NONE), coordinator);

        // assign throws
        KafkaException exc = assertThrows(KafkaException.class, () -> consumer.updateAssignmentMetadataIfNeeded(time.timer(Long.MAX_VALUE)));
        assertEquals(partitionAssigned + singleTopicPartition, exc.getCause().getMessage());

        // the assignment is still updated regardless of the exception
        assertEquals(Set.of(tp0), subscription.assignedPartitions());

        // close's revoke throws
        exc = assertThrows(KafkaException.class, () -> consumer.close(CloseOptions.timeout(Duration.ofMillis(0))));
        assertEquals(partitionRevoked + singleTopicPartition, exc.getCause().getCause().getMessage());

        consumer.close(CloseOptions.timeout(Duration.ofMillis(0)));

        // the assignment is still updated regardless of the exception
        assertTrue(subscription.assignedPartitions().isEmpty());
    }

    // TODO: this test requires rebalance logic which is not yet implemented in the CONSUMER group protocol.
    //       Once it is implemented, this should use both group protocols.
    @ParameterizedTest
    @EnumSource(value = GroupProtocol.class, names = "CLASSIC")
    public void testReturnRecordsDuringRebalance(GroupProtocol groupProtocol) throws InterruptedException {
        Time time = new MockTime(1L);
        ConsumerMetadata metadata = createMetadata(subscription);
        MockClient client = new MockClient(time, metadata);
        ConsumerPartitionAssignor assignor = new CooperativeStickyAssignor();
        KafkaConsumer<String, String> consumer = newConsumer(groupProtocol, time, client, subscription, metadata, assignor, true, groupInstanceId);

        initMetadata(client, Utils.mkMap(Utils.mkEntry(topic, 1), Utils.mkEntry(topic2, 1), Utils.mkEntry(topic3, 1)));

        consumer.subscribe(Arrays.asList(topic, topic2), getConsumerRebalanceListener(consumer));

        Node node = metadata.fetch().nodes().get(0);
        Node coordinator = prepareRebalance(client, node, assignor, Arrays.asList(tp0, t2p0), null);

        // a poll with non-zero milliseconds would complete three round-trips (discover, join, sync)
        TestUtils.waitForCondition(() -> {
            consumer.poll(Duration.ofMillis(100L));
            return consumer.assignment().equals(Set.of(tp0, t2p0));
        }, "Does not complete rebalance in time");

        assertEquals(Set.of(topic, topic2), consumer.subscription());
        assertEquals(Set.of(tp0, t2p0), consumer.assignment());

        // prepare a response of the outstanding fetch so that we have data available on the next poll
        Map<TopicPartition, FetchInfo> fetches1 = new HashMap<>();
        fetches1.put(tp0, new FetchInfo(0, 1));
        fetches1.put(t2p0, new FetchInfo(0, 10));
        client.respondFrom(fetchResponse(fetches1), node);

        ConsumerRecords<String, String> records = consumer.poll(Duration.ZERO);

        // verify that the fetch occurred as expected
        assertEquals(11, records.count());
        assertEquals(1L, consumer.position(tp0));
        assertEquals(10L, consumer.position(t2p0));
        assertEquals(2, records.nextOffsets().size());
        assertEquals(new OffsetAndMetadata(1L), records.nextOffsets().get(tp0));
        assertEquals(new OffsetAndMetadata(10L), records.nextOffsets().get(t2p0));

        // prepare the next response of the prefetch
        fetches1.clear();
        fetches1.put(tp0, new FetchInfo(1, 1));
        fetches1.put(t2p0, new FetchInfo(10, 20));
        client.respondFrom(fetchResponse(fetches1), node);

        // subscription change
        consumer.subscribe(Arrays.asList(topic, topic3), getConsumerRebalanceListener(consumer));

        // verify that subscription has changed but assignment is still unchanged
        assertEquals(Set.of(topic, topic3), consumer.subscription());
        assertEquals(Set.of(tp0, t2p0), consumer.assignment());

        // mock the offset commit response for to be revoked partitions
        Map<TopicPartition, Long> partitionOffsets1 = new HashMap<>();
        partitionOffsets1.put(t2p0, 10L);
        AtomicBoolean commitReceived = prepareOffsetCommitResponse(client, coordinator, partitionOffsets1);

        // poll once which would not complete the rebalance
        records = consumer.poll(Duration.ZERO);

        // clear out the prefetch so it doesn't interfere with the rest of the test
        fetches1.clear();
        fetches1.put(tp0, new FetchInfo(2, 1));
        client.respondFrom(fetchResponse(fetches1), node);

        // verify that the fetch still occurred as expected
        assertEquals(Set.of(topic, topic3), consumer.subscription());
        assertEquals(Set.of(tp0), consumer.assignment());
        assertEquals(1, records.count());
        assertEquals(2L, consumer.position(tp0));
        assertEquals(1, records.nextOffsets().size());
        assertEquals(new OffsetAndMetadata(2L), records.nextOffsets().get(tp0));

        // verify that the offset commits occurred as expected
        assertTrue(commitReceived.get());

        // mock rebalance responses
        client.respondFrom(joinGroupFollowerResponse(assignor, 2, "memberId", "leaderId", Errors.NONE), coordinator);

        // we need to poll 1) for getting the join response, and then send the sync request;
        //                 2) for getting the sync response
        records = consumer.poll(Duration.ZERO);

        // should not finish the response yet
        assertEquals(Set.of(topic, topic3), consumer.subscription());
        assertEquals(Set.of(tp0), consumer.assignment());
        assertEquals(1, records.count());
        assertEquals(3L, consumer.position(tp0));
        assertEquals(1, records.nextOffsets().size());
        assertEquals(new OffsetAndMetadata(3L), records.nextOffsets().get(tp0));

        fetches1.clear();
        fetches1.put(tp0, new FetchInfo(3, 1));
        client.respondFrom(fetchResponse(fetches1), node);

        // now complete the rebalance
        client.respondFrom(syncGroupResponse(Arrays.asList(tp0, t3p0), Errors.NONE), coordinator);

        AtomicInteger count = new AtomicInteger(0);
        AtomicReference<ConsumerRecords<String, String>> recs1 = new AtomicReference<>();
        TestUtils.waitForCondition(() -> {
            recs1.set(consumer.poll(Duration.ofMillis(100L)));
            return consumer.assignment().equals(Set.of(tp0, t3p0)) && count.addAndGet(recs1.get().count()) == 1;

        }, "Does not complete rebalance in time");

        // should have t3 but not sent yet the t3 records
        assertEquals(Set.of(topic, topic3), consumer.subscription());
        assertEquals(Set.of(tp0, t3p0), consumer.assignment());
        assertEquals(4L, consumer.position(tp0));
        assertEquals(0L, consumer.position(t3p0));
        assertEquals(1, recs1.get().nextOffsets().size());
        assertEquals(new OffsetAndMetadata(4L), recs1.get().nextOffsets().get(tp0));

        fetches1.clear();
        fetches1.put(tp0, new FetchInfo(4, 1));
        fetches1.put(t3p0, new FetchInfo(0, 100));
        client.respondFrom(fetchResponse(fetches1), node);

        count.set(0);
        AtomicReference<ConsumerRecords<String, String>> recs2 = new AtomicReference<>();
        TestUtils.waitForCondition(() -> {
            recs2.set(consumer.poll(Duration.ofMillis(100L)));
            return count.addAndGet(recs2.get().count()) == 101;

        }, "Does not complete rebalance in time");

        assertEquals(5L, consumer.position(tp0));
        assertEquals(100L, consumer.position(t3p0));
        assertEquals(2, recs2.get().nextOffsets().size());
        assertEquals(new OffsetAndMetadata(5L), recs2.get().nextOffsets().get(tp0));
        assertEquals(new OffsetAndMetadata(100L), recs2.get().nextOffsets().get(t3p0));

        client.requests().clear();
        consumer.unsubscribe();
        consumer.close(CloseOptions.timeout(Duration.ZERO));
    }

    // TODO: this test requires rebalance logic which is not yet implemented in the CONSUMER group protocol.
    //       Once it is implemented, this should use both group protocols.
    @ParameterizedTest
    @EnumSource(value = GroupProtocol.class, names = "CLASSIC")
    public void testGetGroupMetadata(GroupProtocol groupProtocol) {
        final ConsumerMetadata metadata = createMetadata(subscription);
        final MockClient client = new MockClient(time, metadata);

        initMetadata(client, Map.of(topic, 1));
        final Node node = metadata.fetch().nodes().get(0);

        final KafkaConsumer<String, String> consumer = newConsumer(groupProtocol, time, client, subscription, metadata, assignor, true, groupInstanceId);

        final ConsumerGroupMetadata groupMetadataOnStart = consumer.groupMetadata();
        assertEquals(groupId, groupMetadataOnStart.groupId());
        assertEquals(JoinGroupRequest.UNKNOWN_MEMBER_ID, groupMetadataOnStart.memberId());
        assertEquals(JoinGroupRequest.UNKNOWN_GENERATION_ID, groupMetadataOnStart.generationId());
        assertEquals(groupInstanceId, groupMetadataOnStart.groupInstanceId());

        consumer.subscribe(Set.of(topic), getConsumerRebalanceListener(consumer));
        prepareRebalance(client, node, assignor, List.of(tp0), null);

        // initial fetch
        client.prepareResponseFrom(fetchResponse(tp0, 0, 0), node);
        consumer.updateAssignmentMetadataIfNeeded(time.timer(Long.MAX_VALUE));

        final ConsumerGroupMetadata groupMetadataAfterPoll = consumer.groupMetadata();
        assertEquals(groupId, groupMetadataAfterPoll.groupId());
        assertEquals(memberId, groupMetadataAfterPoll.memberId());
        assertEquals(1, groupMetadataAfterPoll.generationId());
        assertEquals(groupInstanceId, groupMetadataAfterPoll.groupInstanceId());
    }

    @ParameterizedTest
    @EnumSource(value = GroupProtocol.class)
    public void testInvalidGroupMetadata(GroupProtocol groupProtocol) throws InterruptedException {
        ConsumerMetadata metadata = createMetadata(subscription);
        MockClient client = new MockClient(time, metadata);
        initMetadata(client, Map.of(topic, 1));
        KafkaConsumer<String, String> consumer = newConsumer(groupProtocol, time, client, subscription, metadata,
                new RoundRobinAssignor(), true, groupInstanceId);
        consumer.subscribe(List.of(topic));
        // concurrent access is illegal
        client.enableBlockingUntilWakeup(1);
        ExecutorService service = Executors.newSingleThreadExecutor();
        service.execute(() -> consumer.poll(Duration.ofSeconds(5)));
        try {
            TimeUnit.SECONDS.sleep(1);
            assertThrows(ConcurrentModificationException.class, consumer::groupMetadata);
            client.wakeup();
            consumer.wakeup();
        } finally {
            service.shutdown();
            assertTrue(service.awaitTermination(10, TimeUnit.SECONDS));
        }

        // accessing closed consumer is illegal
        consumer.close(CloseOptions.timeout(Duration.ZERO));
        assertThrows(IllegalStateException.class, consumer::groupMetadata);
    }

    @ParameterizedTest
    @EnumSource(GroupProtocol.class)
    @SuppressWarnings("unchecked")
    public void testCurrentLag(GroupProtocol groupProtocol) throws InterruptedException {
        final ConsumerMetadata metadata = createMetadata(subscription);
        final MockClient client = new MockClient(time, metadata);

        initMetadata(client, Map.of(topic, 1));

        consumer = newConsumer(groupProtocol, time, client, subscription, metadata, assignor, true, groupInstanceId);

        // throws for unassigned partition
        assertThrows(IllegalStateException.class, () -> consumer.currentLag(tp0));

        consumer.assign(Set.of(tp0));

        // poll once to update with the current metadata
        consumer.poll(Duration.ofMillis(0));
        TestUtils.waitForCondition(() -> requestGenerated(client, ApiKeys.FIND_COORDINATOR),
                "No metadata requests sent");
        client.respond(FindCoordinatorResponse.prepareResponse(Errors.NONE, groupId, metadata.fetch().nodes().get(0)));

        // no error for no current position
        assertEquals(OptionalLong.empty(), consumer.currentLag(tp0));
        if (groupProtocol == GroupProtocol.CLASSIC) {
            // Classic consumer does not send the LIST_OFFSETS right away (requires an explicit poll),
            // different from the new async consumer, that will send the LIST_OFFSETS request in the background thread
            // on the next background thread poll.
            assertEquals(0, client.inFlightRequestCount());
        }
        // poll once again, which should send the list-offset request
        consumer.seek(tp0, 50L);
        // requests: list-offset, fetch
        TestUtils.waitForCondition(() -> {
            consumer.poll(Duration.ofMillis(0));
            boolean hasListOffsetRequest = requestGenerated(client, ApiKeys.LIST_OFFSETS);
            boolean hasFetchRequest = requestGenerated(client, ApiKeys.FETCH);
            return hasListOffsetRequest && hasFetchRequest;
        }, "No list-offset & fetch request sent");

        // no error for no end offset (so unknown lag)
        assertEquals(OptionalLong.empty(), consumer.currentLag(tp0));

        // poll once again, which should return the list-offset response
        // and hence next call would return correct lag result
        ClientRequest listOffsetRequest = findRequest(client, ApiKeys.LIST_OFFSETS);
        client.respondToRequest(listOffsetRequest, listOffsetsResponse(Map.of(tp0, 90L)));
        consumer.poll(Duration.ofMillis(0));

        // For AsyncKafkaConsumer, subscription state is updated in background, so the result will eventually be updated.
        TestUtils.waitForCondition(() -> {
            OptionalLong result = consumer.currentLag(tp0);
            return result.isPresent() && result.getAsLong() == 40L;
        }, "Subscription state is not updated");
        // requests: fetch
        TestUtils.waitForCondition(() -> requestGenerated(client, ApiKeys.FETCH), "No fetch request sent");

        // one successful fetch should update the log end offset and the position
        ClientRequest fetchRequest = findRequest(client, ApiKeys.FETCH);
        final FetchInfo fetchInfo = new FetchInfo(1L, 99L, 50L, 5);
        client.respondToRequest(fetchRequest, fetchResponse(Map.of(tp0, fetchInfo)));

        final ConsumerRecords<String, String> records = (ConsumerRecords<String, String>) consumer.poll(Duration.ofMillis(1));
        assertEquals(5, records.count());
        assertEquals(55L, consumer.position(tp0));
        assertEquals(1, records.nextOffsets().size());
        assertEquals(new OffsetAndMetadata(55L), records.nextOffsets().get(tp0));

        // correct lag result
        assertEquals(OptionalLong.of(45L), consumer.currentLag(tp0));
    }

    @ParameterizedTest
    @EnumSource(GroupProtocol.class)
    public void testListOffsetShouldUpdateSubscriptions(GroupProtocol groupProtocol) {
        final ConsumerMetadata metadata = createMetadata(subscription);
        final MockClient client = new MockClient(time, metadata);

        initMetadata(client, Map.of(topic, 1));

        consumer = newConsumer(groupProtocol, time, client, subscription, metadata, assignor, false,
                null, groupInstanceId, false);
        consumer.assign(Set.of(tp0));
        consumer.seek(tp0, 50L);

        client.prepareResponse(request -> request instanceof ListOffsetsRequest, listOffsetsResponse(Map.of(tp0, 90L)));
        assertEquals(Map.of(tp0, 90L), consumer.endOffsets(Set.of(tp0)));
        // correct lag result should be returned as well
        assertEquals(OptionalLong.of(40L), consumer.currentLag(tp0));
    }

    private ClientRequest findRequest(MockClient client, ApiKeys apiKey) {
        Optional<ClientRequest> request = client.requests().stream().filter(r -> r.requestBuilder().apiKey().equals(apiKey)).findFirst();
        assertTrue(request.isPresent(), "No " + apiKey + " request was submitted to the client");
        return request.get();
    }

    private KafkaConsumer<String, String> consumerWithPendingAuthenticationError(GroupProtocol groupProtocol,
                                                                                 final Time time) {
        ConsumerMetadata metadata = createMetadata(subscription);
        MockClient client = new MockClient(time, metadata);

        initMetadata(client, Map.of(topic, 1));
        Node node = metadata.fetch().nodes().get(0);

        ConsumerPartitionAssignor assignor = new RangeAssignor();

        client.createPendingAuthenticationError(node, 0);
        return newConsumer(groupProtocol, time, client, subscription, metadata, assignor, false, groupInstanceId);
    }

    private KafkaConsumer<String, String> consumerWithPendingAuthenticationError(GroupProtocol groupProtocol) {
        return consumerWithPendingAuthenticationError(groupProtocol, new MockTime());
    }

    private KafkaConsumer<String, String> consumerWithPendingError(GroupProtocol groupProtocol, final Time time) {
        return consumerWithPendingAuthenticationError(groupProtocol, time);
    }

    private ConsumerRebalanceListener getConsumerRebalanceListener(final KafkaConsumer<?, ?> consumer) {
        return new ConsumerRebalanceListener() {
            @Override
            public void onPartitionsRevoked(Collection<TopicPartition> partitions) {
            }

            @Override
            public void onPartitionsAssigned(Collection<TopicPartition> partitions) {
                // set initial position so we don't need a lookup
                for (TopicPartition partition : partitions)
                    consumer.seek(partition, 0);
            }
        };
    }

    private ConsumerRebalanceListener getExceptionConsumerRebalanceListener() {
        return new ConsumerRebalanceListener() {
            @Override
            public void onPartitionsRevoked(Collection<TopicPartition> partitions) {
                throw new RuntimeException(partitionRevoked + partitions);
            }

            @Override
            public void onPartitionsAssigned(Collection<TopicPartition> partitions) {
                throw new RuntimeException(partitionAssigned + partitions);
            }

            @Override
            public void onPartitionsLost(Collection<TopicPartition> partitions) {
                throw new RuntimeException(partitionLost + partitions);
            }
        };
    }

    private ConsumerMetadata createMetadata(SubscriptionState subscription) {
        return new ConsumerMetadata(0, 0, Long.MAX_VALUE, false, false,
                                    subscription, new LogContext(), new ClusterResourceListeners());
    }

    private Node prepareRebalance(MockClient client, Node node, final Set<String> subscribedTopics, ConsumerPartitionAssignor assignor, List<TopicPartition> partitions, Node coordinator) {
        if (coordinator == null) {
            // lookup coordinator
            client.prepareResponseFrom(FindCoordinatorResponse.prepareResponse(Errors.NONE, groupId, node), node);
            coordinator = new Node(Integer.MAX_VALUE - node.id(), node.host(), node.port());
        }

        // join group
        client.prepareResponseFrom(body -> {
            JoinGroupRequest joinGroupRequest = (JoinGroupRequest) body;
            Iterator<JoinGroupRequestData.JoinGroupRequestProtocol> protocolIterator =
                    joinGroupRequest.data().protocols().iterator();
            assertTrue(protocolIterator.hasNext());

            ByteBuffer protocolMetadata = ByteBuffer.wrap(protocolIterator.next().metadata());
            ConsumerPartitionAssignor.Subscription subscription = ConsumerProtocol.deserializeSubscription(protocolMetadata);
            return subscribedTopics.equals(new HashSet<>(subscription.topics()));
        }, joinGroupFollowerResponse(assignor, 1, memberId, leaderId, Errors.NONE), coordinator);

        // sync group
        client.prepareResponseFrom(syncGroupResponse(partitions, Errors.NONE), coordinator);

        return coordinator;
    }

    private Node prepareRebalance(MockClient client, Node node, ConsumerPartitionAssignor assignor, List<TopicPartition> partitions, Node coordinator) {
        if (coordinator == null) {
            // lookup coordinator
            client.prepareResponseFrom(FindCoordinatorResponse.prepareResponse(Errors.NONE, groupId, node), node);
            coordinator = new Node(Integer.MAX_VALUE - node.id(), node.host(), node.port());
        }

        // join group
        client.prepareResponseFrom(joinGroupFollowerResponse(assignor, 1, memberId, leaderId, Errors.NONE), coordinator);

        // sync group
        client.prepareResponseFrom(syncGroupResponse(partitions, Errors.NONE), coordinator);

        return coordinator;
    }

    private AtomicBoolean prepareHeartbeatResponse(MockClient client, Node coordinator, Errors error) {
        final AtomicBoolean heartbeatReceived = new AtomicBoolean(false);
        client.prepareResponseFrom(body -> {
            heartbeatReceived.set(true);
            return true;
        }, new HeartbeatResponse(new HeartbeatResponseData().setErrorCode(error.code())), coordinator);
        return heartbeatReceived;
    }

    private AtomicBoolean prepareOffsetCommitResponse(MockClient client, Node coordinator, final Map<TopicPartition, Long> partitionOffsets) {
        final AtomicBoolean commitReceived = new AtomicBoolean(true);
        Map<TopicPartition, Errors> response = new HashMap<>();
        for (TopicPartition partition : partitionOffsets.keySet())
            response.put(partition, Errors.NONE);

        client.prepareResponseFrom(body -> {
            OffsetCommitRequest commitRequest = (OffsetCommitRequest) body;
            Map<TopicPartition, Long> commitErrors = commitRequest.offsets();

            for (Map.Entry<TopicPartition, Long> partitionOffset : partitionOffsets.entrySet()) {
                // verify that the expected offset has been committed
                if (!commitErrors.get(partitionOffset.getKey()).equals(partitionOffset.getValue())) {
                    commitReceived.set(false);
                    return false;
                }
            }
            return true;
        }, offsetCommitResponse(response), coordinator);
        return commitReceived;
    }

    private AtomicBoolean prepareOffsetCommitResponse(MockClient client, Node coordinator, final TopicPartition partition, final long offset) {
        return prepareOffsetCommitResponse(client, coordinator, Map.of(partition, offset));
    }

    private OffsetCommitResponse offsetCommitResponse(Map<TopicPartition, Errors> responseData) {
        return new OffsetCommitResponse(responseData);
    }

    private JoinGroupResponse joinGroupFollowerResponse(ConsumerPartitionAssignor assignor, int generationId, String memberId, String leaderId, Errors error) {
        return new JoinGroupResponse(
                new JoinGroupResponseData()
                        .setErrorCode(error.code())
                        .setGenerationId(generationId)
                        .setProtocolName(assignor.name())
                        .setLeader(leaderId)
                        .setMemberId(memberId)
                        .setMembers(Collections.emptyList()),
                ApiKeys.JOIN_GROUP.latestVersion()
        );
    }

    private SyncGroupResponse syncGroupResponse(List<TopicPartition> partitions, Errors error) {
        ByteBuffer buf = ConsumerProtocol.serializeAssignment(new ConsumerPartitionAssignor.Assignment(partitions));
        return new SyncGroupResponse(
                new SyncGroupResponseData()
                        .setErrorCode(error.code())
                        .setAssignment(Utils.toArray(buf))
        );
    }

    private OffsetFetchResponse offsetResponse(Map<TopicPartition, Long> offsets, Errors error) {
        var grouped = offsets.entrySet().stream().collect(Collectors.groupingBy(e -> e.getKey().topic()));

        return new OffsetFetchResponse(
            new OffsetFetchResponseData()
                .setGroups(List.of(
                    new OffsetFetchResponseData.OffsetFetchResponseGroup()
                        .setGroupId(groupId)
                        .setTopics(grouped.entrySet().stream().map(entry ->
                            new OffsetFetchResponseData.OffsetFetchResponseTopics()
                                .setName(entry.getKey())
                                .setPartitions(entry.getValue().stream().map(partition ->
                                    new OffsetFetchResponseData.OffsetFetchResponsePartitions()
                                        .setPartitionIndex(partition.getKey().partition())
                                        .setErrorCode(error.code())
                                        .setCommittedOffset(partition.getValue())
                                ).collect(Collectors.toList()))
                        ).collect(Collectors.toList()))
                )),
            ApiKeys.OFFSET_FETCH.latestVersion()
        );
    }

    private ListOffsetsResponse listOffsetsResponse(Map<TopicPartition, Long> offsets) {
        return listOffsetsResponse(offsets, Collections.emptyMap());
    }

    private ListOffsetsResponse listOffsetsResponse(Map<TopicPartition, Long> partitionOffsets,
                                                    Map<TopicPartition, Errors> partitionErrors) {
        Map<String, ListOffsetsTopicResponse> responses = new HashMap<>();
        for (Map.Entry<TopicPartition, Long> partitionOffset : partitionOffsets.entrySet()) {
            TopicPartition tp = partitionOffset.getKey();
            ListOffsetsTopicResponse topic = responses.computeIfAbsent(tp.topic(), k -> new ListOffsetsTopicResponse().setName(tp.topic()));
            topic.partitions().add(new ListOffsetsPartitionResponse()
                    .setPartitionIndex(tp.partition())
                    .setErrorCode(Errors.NONE.code())
                    .setTimestamp(ListOffsetsResponse.UNKNOWN_TIMESTAMP)
                    .setOffset(partitionOffset.getValue()));
        }

        for (Map.Entry<TopicPartition, Errors> partitionError : partitionErrors.entrySet()) {
            TopicPartition tp = partitionError.getKey();
            ListOffsetsTopicResponse topic = responses.computeIfAbsent(tp.topic(), k -> new ListOffsetsTopicResponse().setName(tp.topic()));
            topic.partitions().add(new ListOffsetsPartitionResponse()
                    .setPartitionIndex(tp.partition())
                    .setErrorCode(partitionError.getValue().code())
                    .setTimestamp(ListOffsetsResponse.UNKNOWN_TIMESTAMP)
                    .setOffset(ListOffsetsResponse.UNKNOWN_OFFSET));
        }
        ListOffsetsResponseData data = new ListOffsetsResponseData()
                .setTopics(new ArrayList<>(responses.values()));
        return new ListOffsetsResponse(data);
    }

    private FetchResponse fetchResponse(Map<TopicPartition, FetchInfo> fetches) {
        LinkedHashMap<TopicIdPartition, FetchResponseData.PartitionData> tpResponses = new LinkedHashMap<>();
        for (Map.Entry<TopicPartition, FetchInfo> fetchEntry : fetches.entrySet()) {
            TopicPartition partition = fetchEntry.getKey();
            long fetchOffset = fetchEntry.getValue().offset;
            int fetchCount = fetchEntry.getValue().count;
            final long highWatermark = fetchEntry.getValue().logLastOffset + 1;
            final long logStartOffset = fetchEntry.getValue().logFirstOffset;
            final MemoryRecords records;
            if (fetchCount == 0) {
                records = MemoryRecords.EMPTY;
            } else {
                try (MemoryRecordsBuilder builder = MemoryRecords.builder(ByteBuffer.allocate(1024), Compression.NONE,
                        TimestampType.CREATE_TIME, fetchOffset)) {
                    for (int i = 0; i < fetchCount; i++)
                        builder.append(0L, ("key-" + i).getBytes(), ("value-" + i).getBytes());
                    records = builder.build();
                }
            }
            tpResponses.put(new TopicIdPartition(topicIds.get(partition.topic()), partition),
                new FetchResponseData.PartitionData()
                    .setPartitionIndex(partition.partition())
                    .setHighWatermark(highWatermark)
                    .setLogStartOffset(logStartOffset)
                    .setRecords(records));
        }
        return FetchResponse.of(Errors.NONE, 0, INVALID_SESSION_ID, tpResponses, List.of());
    }

    private FetchResponse fetchResponse(TopicPartition partition, long fetchOffset, int count) {
        FetchInfo fetchInfo = new FetchInfo(fetchOffset, count);
        return fetchResponse(Map.of(partition, fetchInfo));
    }

    private KafkaConsumer<String, String> newConsumer(GroupProtocol groupProtocol,
                                                      Time time,
                                                      KafkaClient client,
                                                      SubscriptionState subscription,
                                                      ConsumerMetadata metadata,
                                                      ConsumerPartitionAssignor assignor,
                                                      boolean autoCommitEnabled,
                                                      Optional<String> groupInstanceId) {
        return newConsumer(
            groupProtocol,
            time,
            client,
            subscription,
            metadata,
            assignor,
            autoCommitEnabled,
            groupId,
            groupInstanceId,
            false
        );
    }

    private KafkaConsumer<String, String> newConsumerNoAutoCommit(GroupProtocol groupProtocol,
                                                                  Time time,
                                                                  KafkaClient client,
                                                                  SubscriptionState subscription,
                                                                  ConsumerMetadata metadata) {
        return newConsumer(
            groupProtocol,
            time,
            client,
            subscription,
            metadata,
            new RangeAssignor(),
            false,
            groupId,
            groupInstanceId,
            false
        );
    }

    private KafkaConsumer<String, String> newConsumer(GroupProtocol groupProtocol,
                                                      Time time,
                                                      KafkaClient client,
                                                      SubscriptionState subscription,
                                                      ConsumerMetadata metadata,
                                                      ConsumerPartitionAssignor assignor,
                                                      boolean autoCommitEnabled,
                                                      String groupId,
                                                      Optional<String> groupInstanceId,
                                                      boolean throwOnStableOffsetNotSupported) {
        return newConsumer(
            groupProtocol,
            time,
            client,
            subscription,
            metadata,
            assignor,
            autoCommitEnabled,
            groupId,
            groupInstanceId,
            Optional.of(new StringDeserializer()),
            throwOnStableOffsetNotSupported
        );
    }

    private KafkaConsumer<String, String> newConsumer(GroupProtocol groupProtocol,
                                                      Time time,
                                                      KafkaClient client,
                                                      SubscriptionState subscriptions,
                                                      ConsumerMetadata metadata,
                                                      ConsumerPartitionAssignor assignor,
                                                      boolean autoCommitEnabled,
                                                      String groupId,
                                                      Optional<String> groupInstanceId,
                                                      Optional<Deserializer<String>> valueDeserializerOpt,
                                                      boolean throwOnStableOffsetNotSupported) {
        Deserializer<String> keyDeserializer = new StringDeserializer();
        Deserializer<String> valueDeserializer = valueDeserializerOpt.orElse(new StringDeserializer());
        LogContext logContext = new LogContext();
        List<ConsumerPartitionAssignor> assignors = List.of(assignor);
        ConsumerConfig config = newConsumerConfig(
            groupProtocol,
            autoCommitEnabled,
            groupId,
            groupInstanceId,
            valueDeserializer,
            throwOnStableOffsetNotSupported
        );
        return new KafkaConsumer<>(
            logContext,
            time,
            config,
            keyDeserializer,
            valueDeserializer,
            client,
            subscriptions,
            metadata,
            assignors
        );
    }

    private ConsumerConfig newConsumerConfig(GroupProtocol groupProtocol,
                                             boolean autoCommitEnabled,
                                             String groupId,
                                             Optional<String> groupInstanceId,
                                             Deserializer<String> valueDeserializer,
                                             boolean throwOnStableOffsetNotSupported) {
        String clientId = "mock-consumer";
        long retryBackoffMs = 100;
        long retryBackoffMaxMs = 1000;
        int minBytes = 1;
        int maxBytes = Integer.MAX_VALUE;
        int maxWaitMs = 500;
        int fetchSize = 1024 * 1024;
        int maxPollRecords = Integer.MAX_VALUE;
        boolean checkCrcs = true;
        int rebalanceTimeoutMs = 60000;
        int requestTimeoutMs = defaultApiTimeoutMs / 2;

        Map<String, Object> configs = new HashMap<>();
        configs.put(ConsumerConfig.AUTO_COMMIT_INTERVAL_MS_CONFIG, autoCommitIntervalMs);
        configs.put(ConsumerConfig.CHECK_CRCS_CONFIG, checkCrcs);
        configs.put(ConsumerConfig.CLIENT_ID_CONFIG, clientId);
        configs.put(ConsumerConfig.CLIENT_RACK_CONFIG, CommonClientConfigs.DEFAULT_CLIENT_RACK);
        configs.put(ConsumerConfig.DEFAULT_API_TIMEOUT_MS_CONFIG, defaultApiTimeoutMs);
        configs.put(ConsumerConfig.ENABLE_AUTO_COMMIT_CONFIG, autoCommitEnabled);
        configs.put(ConsumerConfig.FETCH_MAX_BYTES_CONFIG, maxBytes);
        configs.put(ConsumerConfig.FETCH_MAX_WAIT_MS_CONFIG, maxWaitMs);
        configs.put(ConsumerConfig.FETCH_MIN_BYTES_CONFIG, minBytes);
        configs.put(ConsumerConfig.GROUP_ID_CONFIG, groupId);
        configs.put(ConsumerConfig.GROUP_PROTOCOL_CONFIG, groupProtocol.name());
        if (groupProtocol == GroupProtocol.CLASSIC) {
            configs.put(ConsumerConfig.HEARTBEAT_INTERVAL_MS_CONFIG, heartbeatIntervalMs);
            configs.put(ConsumerConfig.SESSION_TIMEOUT_MS_CONFIG, sessionTimeoutMs);
        }
        configs.put(ConsumerConfig.ISOLATION_LEVEL_CONFIG, IsolationLevel.READ_UNCOMMITTED.toString());
        configs.put(ConsumerConfig.KEY_DESERIALIZER_CLASS_CONFIG, StringDeserializer.class);
        configs.put(ConsumerConfig.MAX_PARTITION_FETCH_BYTES_CONFIG, fetchSize);
        configs.put(ConsumerConfig.MAX_POLL_INTERVAL_MS_CONFIG, rebalanceTimeoutMs);
        configs.put(ConsumerConfig.MAX_POLL_RECORDS_CONFIG, maxPollRecords);
        configs.put(ConsumerConfig.REQUEST_TIMEOUT_MS_CONFIG, requestTimeoutMs);
        configs.put(ConsumerConfig.RETRY_BACKOFF_MAX_MS_CONFIG, retryBackoffMaxMs);
        configs.put(ConsumerConfig.RETRY_BACKOFF_MS_CONFIG, retryBackoffMs);
        configs.put(ConsumerConfig.THROW_ON_FETCH_STABLE_OFFSET_UNSUPPORTED, throwOnStableOffsetNotSupported);
        configs.put(ConsumerConfig.VALUE_DESERIALIZER_CLASS_CONFIG, valueDeserializer.getClass());
        configs.put(ConsumerConfig.BOOTSTRAP_SERVERS_CONFIG, "localhost:9092");
        groupInstanceId.ifPresent(gi -> configs.put(ConsumerConfig.GROUP_INSTANCE_ID_CONFIG, gi));

        return new ConsumerConfig(configs);
    }

    private static class FetchInfo {
        long logFirstOffset;
        long logLastOffset;
        long offset;
        int count;

        FetchInfo(long offset, int count) {
            this(0L, offset + count, offset, count);
        }

        FetchInfo(long logFirstOffset, long logLastOffset, long offset, int count) {
            this.logFirstOffset = logFirstOffset;
            this.logLastOffset = logLastOffset;
            this.offset = offset;
            this.count = count;
        }
    }

    @ParameterizedTest
    @EnumSource(GroupProtocol.class)
    public void testSubscriptionOnInvalidTopic(GroupProtocol groupProtocol) throws InterruptedException {
        ConsumerMetadata metadata = createMetadata(subscription);
        MockClient client = new MockClient(time, metadata);

        initMetadata(client, Map.of(topic, 1));
        Cluster cluster = metadata.fetch();

        String invalidTopicName = "topic abc";  // Invalid topic name due to space

        List<MetadataResponse.TopicMetadata> topicMetadata = new ArrayList<>();
        topicMetadata.add(new MetadataResponse.TopicMetadata(Errors.INVALID_TOPIC_EXCEPTION,
                invalidTopicName, false, Collections.emptyList()));
        MetadataResponse updateResponse = RequestTestUtils.metadataResponse(cluster.nodes(),
                cluster.clusterResource().clusterId(),
                cluster.controller().id(),
                topicMetadata);
        client.prepareMetadataUpdate(updateResponse);

        KafkaConsumer<String, String> consumer = newConsumer(groupProtocol, time, client, subscription, metadata, assignor, true, groupInstanceId);
        consumer.subscribe(Set.of(invalidTopicName), getConsumerRebalanceListener(consumer));

        if (groupProtocol == GroupProtocol.CONSUMER) {
            // New consumer poll(ZERO) needs to wait for the event added by a call to poll, to be processed
            // by the background thread, so it can realize there is invalid topics and then
            // throw the InvalidTopicException
            assertPollEventuallyThrows(consumer, InvalidTopicException.class,
                    "Consumer was not able to update fetch positions on continuous calls with 0 timeout");
        } else {
            assertThrows(InvalidTopicException.class, () -> consumer.poll(Duration.ZERO));
        }
    }

    private static <T extends Throwable> void assertPollEventuallyThrows(KafkaConsumer<?, ?> consumer,
            Class<T> expectedException, String errMsg) throws InterruptedException {
        TestUtils.waitForCondition(() -> {
            try {
                consumer.poll(Duration.ZERO);
                return false;
            } catch (Throwable exception) {
                return expectedException.isInstance(exception);
            }
        }, errMsg);
    }

    @ParameterizedTest
    @EnumSource(GroupProtocol.class)
    public void testPollTimeMetrics(GroupProtocol groupProtocol) {
        ConsumerMetadata metadata = createMetadata(subscription);
        MockClient client = new MockClient(time, metadata);
        initMetadata(client, Map.of(topic, 1));

        KafkaConsumer<String, String> consumer = newConsumer(groupProtocol, time, client, subscription, metadata, assignor, true, groupInstanceId);
        consumer.subscribe(List.of(topic));
        // MetricName objects to check
        Metrics metrics = consumer.metricsRegistry();
        MetricName lastPollSecondsAgoName = metrics.metricName("last-poll-seconds-ago", "consumer-metrics");
        MetricName timeBetweenPollAvgName = metrics.metricName("time-between-poll-avg", "consumer-metrics");
        MetricName timeBetweenPollMaxName = metrics.metricName("time-between-poll-max", "consumer-metrics");
        // Test default values
        assertEquals(-1.0d, consumer.metrics().get(lastPollSecondsAgoName).metricValue());
        assertEquals(Double.NaN, consumer.metrics().get(timeBetweenPollAvgName).metricValue());
        assertEquals(Double.NaN, consumer.metrics().get(timeBetweenPollMaxName).metricValue());
        // Call first poll
        consumer.poll(Duration.ZERO);
        assertEquals(0.0d, consumer.metrics().get(lastPollSecondsAgoName).metricValue());
        assertEquals(0.0d, consumer.metrics().get(timeBetweenPollAvgName).metricValue());
        assertEquals(0.0d, consumer.metrics().get(timeBetweenPollMaxName).metricValue());
        // Advance time by 5,000 (total time = 5,000)
        time.sleep(5 * 1000L);
        assertEquals(5.0d, consumer.metrics().get(lastPollSecondsAgoName).metricValue());
        // Call second poll
        consumer.poll(Duration.ZERO);
        assertEquals(2.5 * 1000d, consumer.metrics().get(timeBetweenPollAvgName).metricValue());
        assertEquals(5 * 1000d, consumer.metrics().get(timeBetweenPollMaxName).metricValue());
        // Advance time by 10,000 (total time = 15,000)
        time.sleep(10 * 1000L);
        assertEquals(10.0d, consumer.metrics().get(lastPollSecondsAgoName).metricValue());
        // Call third poll
        consumer.poll(Duration.ZERO);
        assertEquals(5 * 1000d, consumer.metrics().get(timeBetweenPollAvgName).metricValue());
        assertEquals(10 * 1000d, consumer.metrics().get(timeBetweenPollMaxName).metricValue());
        // Advance time by 5,000 (total time = 20,000)
        time.sleep(5 * 1000L);
        assertEquals(5.0d, consumer.metrics().get(lastPollSecondsAgoName).metricValue());
        // Call fourth poll
        consumer.poll(Duration.ZERO);
        assertEquals(5 * 1000d, consumer.metrics().get(timeBetweenPollAvgName).metricValue());
        assertEquals(10 * 1000d, consumer.metrics().get(timeBetweenPollMaxName).metricValue());
    }

    @ParameterizedTest
    @EnumSource(GroupProtocol.class)
public void testPollIdleRatio(GroupProtocol groupProtocol) {
        ConsumerMetadata metadata = createMetadata(subscription);
        MockClient client = new MockClient(time, metadata);
        initMetadata(client, Map.of(topic, 1));

        KafkaConsumer<String, String> consumer = newConsumer(groupProtocol, time, client, subscription, metadata, assignor, true, groupInstanceId);
        // MetricName object to check
        Metrics metrics = consumer.metricsRegistry();
        MetricName pollIdleRatio = metrics.metricName("poll-idle-ratio-avg", "consumer-metrics");
        // Test default value
        assertEquals(Double.NaN, consumer.metrics().get(pollIdleRatio).metricValue());

        // 1st poll
        // Spend 50ms in poll so value = 1.0
        consumer.kafkaConsumerMetrics().recordPollStart(time.milliseconds());
        time.sleep(50);
        consumer.kafkaConsumerMetrics().recordPollEnd(time.milliseconds());

        assertEquals(1.0d, consumer.metrics().get(pollIdleRatio).metricValue());

        // 2nd poll
        // Spend 50m outside poll and 0ms in poll so value = 0.0
        time.sleep(50);
        consumer.kafkaConsumerMetrics().recordPollStart(time.milliseconds());
        consumer.kafkaConsumerMetrics().recordPollEnd(time.milliseconds());

        // Avg of first two data points
        assertEquals((1.0d + 0.0d) / 2, consumer.metrics().get(pollIdleRatio).metricValue());

        // 3rd poll
        // Spend 25ms outside poll and 25ms in poll so value = 0.5
        time.sleep(25);
        consumer.kafkaConsumerMetrics().recordPollStart(time.milliseconds());
        time.sleep(25);
        consumer.kafkaConsumerMetrics().recordPollEnd(time.milliseconds());

        // Avg of three data points
        assertEquals((1.0d + 0.0d + 0.5d) / 3, consumer.metrics().get(pollIdleRatio).metricValue());
    }

    private static boolean consumerMetricPresent(KafkaConsumer<String, String> consumer, String name) {
        MetricName metricName = new MetricName(name, "consumer-metrics", "", Collections.emptyMap());
        return consumer.metricsRegistry().metrics().containsKey(metricName);
    }

    @ParameterizedTest
    @EnumSource(GroupProtocol.class)
    public void testClosingConsumerUnregistersConsumerMetrics(GroupProtocol groupProtocol) {
        Time time = new MockTime(1L);
        ConsumerMetadata metadata = createMetadata(subscription);
        MockClient client = new MockClient(time, metadata);
        initMetadata(client, Map.of(topic, 1));
        KafkaConsumer<String, String> consumer = newConsumer(groupProtocol, time, client, subscription, metadata,
            new RoundRobinAssignor(), true, groupInstanceId);
        consumer.subscribe(List.of(topic));
        assertTrue(consumerMetricPresent(consumer, "last-poll-seconds-ago"));
        assertTrue(consumerMetricPresent(consumer, "time-between-poll-avg"));
        assertTrue(consumerMetricPresent(consumer, "time-between-poll-max"));
        consumer.close();
        assertFalse(consumerMetricPresent(consumer, "last-poll-seconds-ago"));
        assertFalse(consumerMetricPresent(consumer, "time-between-poll-avg"));
        assertFalse(consumerMetricPresent(consumer, "time-between-poll-max"));
    }

    // NOTE: this test uses the enforceRebalance API which is not implemented in the CONSUMER group protocol.
    @ParameterizedTest
    @EnumSource(value = GroupProtocol.class, names = "CLASSIC")
    public void testEnforceRebalanceWithManualAssignment(GroupProtocol groupProtocol) {
        consumer = newConsumer(groupProtocol, null);
        consumer.assign(Set.of(new TopicPartition("topic", 0)));
        assertThrows(IllegalStateException.class, consumer::enforceRebalance);
    }

    // NOTE: this test uses the enforceRebalance API which is not implemented in the CONSUMER group protocol.
    @ParameterizedTest
    @EnumSource(value = GroupProtocol.class, names = "CLASSIC")
    public void testEnforceRebalanceTriggersRebalanceOnNextPoll(GroupProtocol groupProtocol) {
        Time time = new MockTime(1L);
        ConsumerMetadata metadata = createMetadata(subscription);
        MockClient client = new MockClient(time, metadata);
        KafkaConsumer<String, String> consumer = newConsumer(groupProtocol, time, client, subscription, metadata, assignor, true, groupInstanceId);
        MockRebalanceListener countingRebalanceListener = new MockRebalanceListener();
        initMetadata(client, Utils.mkMap(Utils.mkEntry(topic, 1), Utils.mkEntry(topic2, 1), Utils.mkEntry(topic3, 1)));

        consumer.subscribe(Arrays.asList(topic, topic2), countingRebalanceListener);
        Node node = metadata.fetch().nodes().get(0);
        prepareRebalance(client, node, assignor, Arrays.asList(tp0, t2p0), null);

        // a first rebalance to get the assignment, we need two poll calls since we need two round trips to finish join / sync-group
        consumer.poll(Duration.ZERO);
        consumer.poll(Duration.ZERO);

        // onPartitionsRevoked is not invoked when first joining the group
        assertEquals(0, countingRebalanceListener.revokedCount);
        assertEquals(1, countingRebalanceListener.assignedCount);

        consumer.enforceRebalance();

        // the next poll should trigger a rebalance
        consumer.poll(Duration.ZERO);

        assertEquals(1, countingRebalanceListener.revokedCount);
    }

    // NOTE: this test uses the enforceRebalance API which is not implemented in the CONSUMER group protocol.
    @ParameterizedTest
    @EnumSource(value = GroupProtocol.class, names = "CLASSIC")
    public void testEnforceRebalanceReason(GroupProtocol groupProtocol) {
        Time time = new MockTime(1L);

        ConsumerMetadata metadata = createMetadata(subscription);
        MockClient client = new MockClient(time, metadata);
        initMetadata(client, Utils.mkMap(Utils.mkEntry(topic, 1)));
        Node node = metadata.fetch().nodes().get(0);

        consumer = newConsumer(
            groupProtocol,
            time,
            client,
            subscription,
            metadata,
            assignor,
            true,
            groupInstanceId
        );
        consumer.subscribe(List.of(topic));

        // Lookup coordinator.
        client.prepareResponseFrom(FindCoordinatorResponse.prepareResponse(Errors.NONE, groupId, node), node);
        consumer.poll(Duration.ZERO);

        // Initial join sends an empty reason.
        prepareJoinGroupAndVerifyReason(client, node, "");
        consumer.poll(Duration.ZERO);

        // A null reason should be replaced by the default reason.
        consumer.enforceRebalance(null);
        prepareJoinGroupAndVerifyReason(client, node, DEFAULT_REASON);
        consumer.poll(Duration.ZERO);

        // An empty reason should be replaced by the default reason.
        consumer.enforceRebalance("");
        prepareJoinGroupAndVerifyReason(client, node, DEFAULT_REASON);
        consumer.poll(Duration.ZERO);

        // A non-null and non-empty reason is sent as-is.
        String customReason = "user provided reason";
        consumer.enforceRebalance(customReason);
        prepareJoinGroupAndVerifyReason(client, node, customReason);
        consumer.poll(Duration.ZERO);
    }

    private void prepareJoinGroupAndVerifyReason(
        MockClient client,
        Node node,
        String expectedReason
    ) {
        client.prepareResponseFrom(
            body -> {
                JoinGroupRequest joinGroupRequest = (JoinGroupRequest) body;
                return expectedReason.equals(joinGroupRequest.data().reason());
            },
            joinGroupFollowerResponse(assignor, 1, memberId, leaderId, Errors.NONE),
            node
        );
    }

    // TODO: this test triggers a bug with the CONSUMER group protocol implementation.
    //       The bug will be investigated and fixed so this test can use both group protocols.
    @ParameterizedTest
    @EnumSource(value = GroupProtocol.class)
    public void configurableObjectsShouldSeeGeneratedClientId(GroupProtocol groupProtocol) {
        CLIENT_IDS.clear();
        Properties props = new Properties();
        props.put(ConsumerConfig.GROUP_PROTOCOL_CONFIG, groupProtocol.name());
        props.put(ConsumerConfig.BOOTSTRAP_SERVERS_CONFIG, "localhost:9999");
        props.put(ConsumerConfig.KEY_DESERIALIZER_CLASS_CONFIG, DeserializerForClientId.class.getName());
        props.put(ConsumerConfig.VALUE_DESERIALIZER_CLASS_CONFIG, DeserializerForClientId.class.getName());
        props.put(ConsumerConfig.INTERCEPTOR_CLASSES_CONFIG, ConsumerInterceptorForClientId.class.getName());

        consumer = newConsumer(props);
        assertNotNull(consumer.clientId());
        assertNotEquals(0, consumer.clientId().length());
        assertEquals(3, CLIENT_IDS.size());
        CLIENT_IDS.forEach(id -> assertEquals(id, consumer.clientId()));
    }

    @ParameterizedTest
    @EnumSource(GroupProtocol.class)
    public void testUnusedConfigs(GroupProtocol groupProtocol) {
        Map<String, Object> props = new HashMap<>();
        props.put(ConsumerConfig.GROUP_PROTOCOL_CONFIG, groupProtocol.name());
        props.put(ConsumerConfig.BOOTSTRAP_SERVERS_CONFIG, "localhost:9999");
        props.put(SslConfigs.SSL_PROTOCOL_CONFIG, "TLS");
        ConsumerConfig config = new ConsumerConfig(ConsumerConfig.appendDeserializerToConfig(props, new StringDeserializer(), new StringDeserializer()));

        assertTrue(config.unused().contains(SslConfigs.SSL_PROTOCOL_CONFIG));

        consumer = new KafkaConsumer<>(config, null, null);
        assertTrue(config.unused().contains(SslConfigs.SSL_PROTOCOL_CONFIG));
    }

    @ParameterizedTest
    @EnumSource(value = GroupProtocol.class, names = "CLASSIC")
    public void testAssignorNameConflict(GroupProtocol groupProtocol) {
        Map<String, Object> configs = new HashMap<>();
        configs.put(ConsumerConfig.GROUP_PROTOCOL_CONFIG, groupProtocol.name());
        configs.put(ConsumerConfig.BOOTSTRAP_SERVERS_CONFIG, "localhost:9999");
        configs.put(ConsumerConfig.PARTITION_ASSIGNMENT_STRATEGY_CONFIG,
            Arrays.asList(RangeAssignor.class.getName(), ConsumerPartitionAssignorTest.TestConsumerPartitionAssignor.class.getName()));

        assertThrows(KafkaException.class,
            () -> newConsumer(configs, new StringDeserializer(), new StringDeserializer()));
    }

    @ParameterizedTest
    @EnumSource(value = GroupProtocol.class)
    public void testOffsetsForTimesTimeout(GroupProtocol groupProtocol) {
        final KafkaConsumer<String, String> consumer = consumerForCheckingTimeoutException(groupProtocol);
        assertEquals(
            "Failed to get offsets by times in 60000ms",
            assertThrows(org.apache.kafka.common.errors.TimeoutException.class, () -> consumer.offsetsForTimes(Map.of(tp0, 0L))).getMessage()
        );
        consumer.close(CloseOptions.timeout(Duration.ZERO));
    }

    @ParameterizedTest
    @EnumSource(value = GroupProtocol.class)
    public void testBeginningOffsetsTimeout(GroupProtocol groupProtocol) {
        final KafkaConsumer<String, String> consumer = consumerForCheckingTimeoutException(groupProtocol);
        assertEquals(
            "Failed to get offsets by times in 60000ms",
            assertThrows(org.apache.kafka.common.errors.TimeoutException.class, () -> consumer.beginningOffsets(List.of(tp0))).getMessage()
        );
        consumer.close(CloseOptions.timeout(Duration.ZERO));
    }

    @ParameterizedTest
    @EnumSource(value = GroupProtocol.class)
    public void testEndOffsetsTimeout(GroupProtocol groupProtocol) {
        final KafkaConsumer<String, String> consumer = consumerForCheckingTimeoutException(groupProtocol);
        assertEquals(
            "Failed to get offsets by times in 60000ms",
            assertThrows(org.apache.kafka.common.errors.TimeoutException.class, () -> consumer.endOffsets(List.of(tp0))).getMessage()
        );
        consumer.close(CloseOptions.timeout(Duration.ZERO));
    }

    @ParameterizedTest
    @EnumSource(GroupProtocol.class)
    public void testClientInstanceId(GroupProtocol groupProtocol) {
        Properties props = new Properties();
        props.setProperty(ConsumerConfig.BOOTSTRAP_SERVERS_CONFIG, "localhost:9999");
        props.setProperty(ConsumerConfig.GROUP_PROTOCOL_CONFIG, groupProtocol.name());

        ClientTelemetryReporter clientTelemetryReporter = mock(ClientTelemetryReporter.class);
        clientTelemetryReporter.configure(any());

        try (MockedStatic<CommonClientConfigs> mockedCommonClientConfigs = mockStatic(CommonClientConfigs.class, new CallsRealMethods())) {
            mockedCommonClientConfigs.when(() -> CommonClientConfigs.telemetryReporter(anyString(), any())).thenReturn(Optional.of(clientTelemetryReporter));

            ClientTelemetrySender clientTelemetrySender = mock(ClientTelemetrySender.class);
            Uuid expectedUuid = Uuid.randomUuid();
            when(clientTelemetryReporter.telemetrySender()).thenReturn(clientTelemetrySender);
            when(clientTelemetrySender.clientInstanceId(any())).thenReturn(Optional.of(expectedUuid));

            consumer = newConsumer(props, new StringDeserializer(), new StringDeserializer());
            Uuid uuid = consumer.clientInstanceId(Duration.ofMillis(0));
            assertEquals(expectedUuid, uuid);
        }
    }

    @ParameterizedTest
    @EnumSource(GroupProtocol.class)
    public void testClientInstanceIdInvalidTimeout(GroupProtocol groupProtocol) {
        Properties props = new Properties();
        props.setProperty(ConsumerConfig.BOOTSTRAP_SERVERS_CONFIG, "localhost:9999");
        props.setProperty(ConsumerConfig.GROUP_PROTOCOL_CONFIG, groupProtocol.name());


        consumer = newConsumer(props, new StringDeserializer(), new StringDeserializer());
        Exception exception = assertThrows(IllegalArgumentException.class, () -> consumer.clientInstanceId(Duration.ofMillis(-1)));
        assertEquals("The timeout cannot be negative.", exception.getMessage());
    }

    @ParameterizedTest
    @EnumSource(GroupProtocol.class)
    public void testClientInstanceIdNoTelemetryReporterRegistered(GroupProtocol groupProtocol) {
        Properties props = new Properties();
        props.setProperty(ConsumerConfig.BOOTSTRAP_SERVERS_CONFIG, "localhost:9999");
        props.setProperty(ConsumerConfig.ENABLE_METRICS_PUSH_CONFIG, "false");
        props.setProperty(ConsumerConfig.GROUP_PROTOCOL_CONFIG, groupProtocol.name());

        consumer = newConsumer(props, new StringDeserializer(), new StringDeserializer());
        Exception exception = assertThrows(IllegalStateException.class, () -> consumer.clientInstanceId(Duration.ofMillis(0)));
        assertEquals("Telemetry is not enabled. Set config `enable.metrics.push` to `true`.", exception.getMessage());
    }

    @SuppressWarnings("resource")
    @ParameterizedTest
    @EnumSource(GroupProtocol.class)
    public void testEmptyGroupId(GroupProtocol groupProtocol) {
        KafkaException e = assertThrows(KafkaException.class, () -> newConsumer(groupProtocol, ""));
        assertInstanceOf(InvalidGroupIdException.class, e.getCause());
        assertEquals("The configured group.id should not be an empty string or whitespace.", e.getCause().getMessage());
    }

    @SuppressWarnings("resource")
    @ParameterizedTest
    @EnumSource(GroupProtocol.class)
    public void testGroupIdWithWhitespace(GroupProtocol groupProtocol) {
        KafkaException e = assertThrows(KafkaException.class, () -> newConsumer(groupProtocol, " "));
        assertInstanceOf(InvalidGroupIdException.class, e.getCause());
        assertEquals("The configured group.id should not be an empty string or whitespace.", e.getCause().getMessage());
    }

    private KafkaConsumer<String, String> consumerForCheckingTimeoutException(GroupProtocol groupProtocol) {
        ConsumerMetadata metadata = createMetadata(subscription);
        MockClient client = new MockClient(time, metadata);

        initMetadata(client, Map.of(topic, 1));

        ConsumerPartitionAssignor assignor = new RangeAssignor();

        if (groupProtocol == GroupProtocol.CONSUMER) {
            Node node = metadata.fetch().nodes().get(0);
            client.prepareResponseFrom(FindCoordinatorResponse.prepareResponse(Errors.NONE, groupId, node), node);
        }

        final KafkaConsumer<String, String> consumer = newConsumer(groupProtocol, time, client, subscription, metadata, assignor, false, groupInstanceId);

        int maxPreparedResponses = GroupProtocol.CLASSIC.equals(groupProtocol) ? 10 : 1;
        for (int i = 0; i < maxPreparedResponses; i++) {
            client.prepareResponse(
                request -> {
                    time.sleep(defaultApiTimeoutMs / maxPreparedResponses);
                    return request instanceof ListOffsetsRequest;
                },
                listOffsetsResponse(
                    Collections.emptyMap(),
                    Map.of(tp0, Errors.UNKNOWN_TOPIC_OR_PARTITION)
                ));
        }

        return consumer;
    }

    @ParameterizedTest
    @EnumSource(GroupProtocol.class)
    public void testCommittedThrowsTimeoutExceptionForNoResponse(GroupProtocol groupProtocol) {
        Time time = new MockTime(Duration.ofSeconds(1).toMillis());

        ConsumerMetadata metadata = createMetadata(subscription);
        MockClient client = new MockClient(time, metadata);

        initMetadata(client, Map.of(topic, 2));
        Node node = metadata.fetch().nodes().get(0);

        client.prepareResponseFrom(FindCoordinatorResponse.prepareResponse(Errors.NONE, groupId, node), node);
        consumer = newConsumer(groupProtocol, time, client, subscription, metadata, assignor, true, groupInstanceId);
        consumer.assign(List.of(tp0));

        // lookup coordinator
        Node coordinator = new Node(Integer.MAX_VALUE - node.id(), node.host(), node.port());

        // try to get committed offsets for one topic-partition - but it is disconnected so there's no response and it will time out
        client.prepareResponseFrom(offsetResponse(Map.of(tp0, 0L), Errors.NONE), coordinator, true);
        org.apache.kafka.common.errors.TimeoutException timeoutException = assertThrows(org.apache.kafka.common.errors.TimeoutException.class,
            () -> consumer.committed(Set.of(tp0), Duration.ofMillis(1000L)));
        assertEquals("Timeout of 1000ms expired before the last committed offset for partitions [test-0] could be determined. " +
            "Try tuning default.api.timeout.ms larger to relax the threshold.", timeoutException.getMessage());
    }

    @ParameterizedTest
    @EnumSource(value = GroupProtocol.class)
    public void testPreventMultiThread(GroupProtocol groupProtocol) throws InterruptedException {
        ConsumerMetadata metadata = createMetadata(subscription);
        MockClient client = new MockClient(time, metadata);
        initMetadata(client, Map.of(topic, 1));
        KafkaConsumer<String, String> consumer = newConsumer(groupProtocol, time, client, subscription, metadata,
                new RoundRobinAssignor(), true, groupInstanceId);
        consumer.subscribe(List.of(topic));


        client.enableBlockingUntilWakeup(1);
        ExecutorService service = Executors.newSingleThreadExecutor();
        service.execute(() -> consumer.poll(Duration.ofSeconds(5)));
        try {
            TimeUnit.SECONDS.sleep(1);
            assertThrows(ConcurrentModificationException.class, () -> consumer.poll(Duration.ZERO));
            client.wakeup();
            consumer.wakeup();
        } finally {
            service.shutdown();
            assertTrue(service.awaitTermination(10, TimeUnit.SECONDS));
        }
    }

    @ParameterizedTest
    @EnumSource(value = GroupProtocol.class)
    public void testPollSendsRequestToJoin(GroupProtocol groupProtocol) throws InterruptedException {
        ConsumerMetadata metadata = createMetadata(subscription);
        MockClient client = new MockClient(time, metadata);
        initMetadata(client, Map.of(topic, 1));
        Node node = metadata.fetch().nodes().get(0);

        client.prepareResponseFrom(FindCoordinatorResponse.prepareResponse(Errors.NONE, groupId, node), node);
        KafkaConsumer<String, String> consumer = newConsumer(groupProtocol, time, client, subscription, metadata,
                assignor, true, groupInstanceId);
        consumer.subscribe(List.of(topic));
        assertFalse(groupProtocol == GroupProtocol.CLASSIC ?
                        requestGenerated(client, ApiKeys.JOIN_GROUP) :
                        requestGenerated(client, ApiKeys.CONSUMER_GROUP_HEARTBEAT),
                "KafkaConsumer#subscribe should not send " + (groupProtocol == GroupProtocol.CLASSIC ? "JoinGroup" : "Heartbeat") + " request");

        consumer.poll(Duration.ZERO);
        TestUtils.waitForCondition(() -> groupProtocol == GroupProtocol.CLASSIC ?
                requestGenerated(client, ApiKeys.JOIN_GROUP) :
                requestGenerated(client, ApiKeys.CONSUMER_GROUP_HEARTBEAT),
                "Expected " + (groupProtocol == GroupProtocol.CLASSIC ? "JoinGroup" : "Heartbeat") + " request");
    }

    @ParameterizedTest
    @EnumSource(value = GroupProtocol.class, names = "CLASSIC")
    public void testSubscribeToRe2jPatternNotSupportedForClassicConsumer(GroupProtocol groupProtocol) {
        KafkaConsumer<String, String> consumer = newConsumerNoAutoCommit(groupProtocol, time, mock(NetworkClient.class), subscription,
            mock(ConsumerMetadata.class));
        assertThrows(UnsupportedOperationException.class, () ->
            consumer.subscribe(new SubscriptionPattern("t*")));
        assertThrows(UnsupportedOperationException.class, () ->
            consumer.subscribe(new SubscriptionPattern("t*"), mock(ConsumerRebalanceListener.class)));
    }

    private boolean requestGenerated(MockClient client, ApiKeys apiKey) {
        return client.requests().stream().anyMatch(request -> request.requestBuilder().apiKey().equals(apiKey));
    }

    private Map<MetricName, KafkaMetric> customMetrics() {
        MetricConfig metricConfig = new MetricConfig();
        Object lock = new Object();
        MetricName metricNameOne = new MetricName("metricOne", "stream-metrics", "description for metric one", new HashMap<>());
        MetricName metricNameTwo = new MetricName("metricTwo", "stream-metrics", "description for metric two", new HashMap<>());

        KafkaMetric streamClientMetricOne = new KafkaMetric(lock, metricNameOne, (Measurable) (m, now) -> 1.0, metricConfig, Time.SYSTEM);
        KafkaMetric streamClientMetricTwo = new KafkaMetric(lock, metricNameTwo, (Measurable) (m, now) -> 2.0, metricConfig, Time.SYSTEM);
        return Map.of(metricNameOne, streamClientMetricOne, metricNameTwo, streamClientMetricTwo);
    }

    private static final List<String> CLIENT_IDS = new ArrayList<>();
    public static class DeserializerForClientId implements Deserializer<byte[]> {
        @Override
        public void configure(Map<String, ?> configs, boolean isKey) {
            CLIENT_IDS.add(configs.get(ConsumerConfig.CLIENT_ID_CONFIG).toString());
        }

        @Override
        public byte[] deserialize(String topic, byte[] data) {
            return data;
        }
    }

    public static class ConsumerInterceptorForClientId implements ConsumerInterceptor<byte[], byte[]> {

        @Override
        public ConsumerRecords<byte[], byte[]> onConsume(ConsumerRecords<byte[], byte[]> records) {
            return records;
        }

        @Override
        public void onCommit(Map<TopicPartition, OffsetAndMetadata> offsets) {

        }

        @Override
        public void close() {

        }

        @Override
        public void configure(Map<String, ?> configs) {
            CLIENT_IDS.add(configs.get(ConsumerConfig.CLIENT_ID_CONFIG).toString());
        }
    }

    @ParameterizedTest
    @EnumSource(value = GroupProtocol.class)
    void testMonitorablePlugins(GroupProtocol groupProtocol) {
        try {
            String clientId = "testMonitorablePlugins";
            Map<String, Object> configs = new HashMap<>();
            configs.put(ConsumerConfig.CLIENT_ID_CONFIG, clientId);
            configs.put(ConsumerConfig.BOOTSTRAP_SERVERS_CONFIG, "localhost:9999");
            configs.put(ConsumerConfig.KEY_DESERIALIZER_CLASS_CONFIG, MonitorableDeserializer.class.getName());
            configs.put(ConsumerConfig.VALUE_DESERIALIZER_CLASS_CONFIG, MonitorableDeserializer.class.getName());
            configs.put(ConsumerConfig.GROUP_PROTOCOL_CONFIG, groupProtocol.name);
            configs.put(ConsumerConfig.INTERCEPTOR_CLASSES_CONFIG, MonitorableInterceptor.class.getName());

            KafkaConsumer<String, String> consumer = new KafkaConsumer<>(configs);
            Map<MetricName, ? extends Metric> metrics = consumer.metrics();

            MetricName expectedKeyDeserializerMetric = expectedMetricName(
                    clientId,
                    ConsumerConfig.KEY_DESERIALIZER_CLASS_CONFIG,
                    MonitorableDeserializer.class);
            assertTrue(metrics.containsKey(expectedKeyDeserializerMetric));
            assertEquals(VALUE, metrics.get(expectedKeyDeserializerMetric).metricValue());

            MetricName expectedValueDeserializerMetric = expectedMetricName(
                    clientId,
                    ConsumerConfig.VALUE_DESERIALIZER_CLASS_CONFIG,
                    MonitorableDeserializer.class);
            assertTrue(metrics.containsKey(expectedValueDeserializerMetric));
            assertEquals(VALUE, metrics.get(expectedValueDeserializerMetric).metricValue());

            MetricName expectedInterceptorMetric = expectedMetricName(
                    clientId,
                    ConsumerConfig.INTERCEPTOR_CLASSES_CONFIG,
                    MonitorableInterceptor.class);
            assertTrue(metrics.containsKey(expectedInterceptorMetric));
            assertEquals(VALUE, metrics.get(expectedInterceptorMetric).metricValue());

            consumer.close(CloseOptions.timeout(Duration.ZERO));
            metrics = consumer.metrics();
            assertFalse(metrics.containsKey(expectedKeyDeserializerMetric));
            assertFalse(metrics.containsKey(expectedValueDeserializerMetric));
            assertFalse(metrics.containsKey(expectedInterceptorMetric));
        } finally {
            MockConsumerInterceptor.resetCounters();
        }
    }

    /**
     * This test ensures that both {@link Consumer} implementations fail on creation when the underlying
     * {@link NetworkClient} fails creation.
     *
     * The logic to check for this case is admittedly a bit awkward because the constructor can fail for all
     * manner of reasons. So a failure case is created by specifying an invalid
     * {@link javax.security.auth.spi.LoginModule} class name, which in turn causes the {@link NetworkClient}
     * to fail.
     *
     * This test was created to validate the change for KAFKA-19394 for the {@link AsyncKafkaConsumer}. The fix
     * should handle the case where failure during initialization of resources (in this test, the underlying
     * {@link NetworkClient}) will not cause the creation of the {@link AsyncKafkaConsumer} to hang.
     */
    @ParameterizedTest
    @EnumSource(value = GroupProtocol.class)
    public void testConstructorFailsOnNetworkClientConstructorFailure(GroupProtocol groupProtocol) {
        Map<String, Object> configs = Map.of(
            ConsumerConfig.KEY_DESERIALIZER_CLASS_CONFIG, StringDeserializer.class.getName(),
            ConsumerConfig.VALUE_DESERIALIZER_CLASS_CONFIG, StringDeserializer.class.getName(),
            CommonClientConfigs.BOOTSTRAP_SERVERS_CONFIG, "localhost:9999",
            CommonClientConfigs.SECURITY_PROTOCOL_CONFIG, SecurityProtocol.SASL_PLAINTEXT.name,
            SaslConfigs.SASL_MECHANISM, "PLAIN",
            SaslConfigs.SASL_JAAS_CONFIG, "org.example.InvalidLoginModule required ;",
            ConsumerConfig.GROUP_PROTOCOL_CONFIG, groupProtocol.name().toLowerCase(Locale.ROOT)
        );

        KafkaException e = assertThrows(KafkaException.class, () -> {
            try (KafkaConsumer<String, String> ignored = new KafkaConsumer<>(configs)) {
                fail("Should not be able to create the consumer");
            }
        });

        assertEquals("Failed to construct kafka consumer", e.getMessage());

        // The root cause is multiple exceptions deep. This code is more concise and should hopefully be trivial
        // to update should the underlying implementation change.
        Throwable cause = e.getCause();
        assertNotNull(cause);
        assertInstanceOf(KafkaException.class, cause);
        assertEquals("Failed to create new NetworkClient", cause.getMessage());

        cause = cause.getCause();
        assertNotNull(cause);
        assertInstanceOf(KafkaException.class, cause);
        assertEquals(LoginException.class.getName() + ": No LoginModule found for org.example.InvalidLoginModule", cause.getMessage());

        cause = cause.getCause();
        assertNotNull(cause);
        assertInstanceOf(LoginException.class, cause);
        assertEquals("No LoginModule found for org.example.InvalidLoginModule", cause.getMessage());
    }

    private MetricName expectedMetricName(String clientId, String config, Class<?> clazz) {
        Map<String, String> expectedTags = new LinkedHashMap<>();
        expectedTags.put("client-id", clientId);
        expectedTags.put("config", config);
        expectedTags.put("class", clazz.getSimpleName());
        expectedTags.putAll(TAGS);
        return new MetricName(NAME, "plugins", DESCRIPTION, expectedTags);
    }

    private static final String NAME = "name";
    private static final String DESCRIPTION = "description";
    private static final LinkedHashMap<String, String> TAGS = new LinkedHashMap<>();
    private static final double VALUE = 123.0;

    static {
        TAGS.put("t1", "v1");
    }

    public static class MonitorableDeserializer extends MockDeserializer implements Monitorable {

        @Override
        public void withPluginMetrics(PluginMetrics metrics) {
            MetricName name = metrics.metricName(NAME, DESCRIPTION, TAGS);
            metrics.addMetric(name, (Measurable) (config, now) -> VALUE);
        }
    }

    public static class MonitorableInterceptor extends MockConsumerInterceptor implements Monitorable {

        @Override
        public void withPluginMetrics(PluginMetrics metrics) {
            MetricName name = metrics.metricName(NAME, DESCRIPTION, TAGS);
            metrics.addMetric(name, (Measurable) (config, now) -> VALUE);
        }
    }

    public static class CloseInterceptor implements ConsumerInterceptor<String, String> {

        public static final AtomicInteger CLOSE_COUNT = new AtomicInteger(0);

        @Override
        public ConsumerRecords<String, String> onConsume(ConsumerRecords<String, String> records) {
            return null;
        }

        @Override
        public void onCommit(Map<TopicPartition, OffsetAndMetadata> offsets) {
            // no-op
        }

        @Override
        public void close() {
            CLOSE_COUNT.incrementAndGet();
        }

        @Override
        public void configure(Map<String, ?> configs) {
            // no-op
        }

        public static void resetCounters() {
            CLOSE_COUNT.set(0);
        }
    }
}
