content: /*
   * 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.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.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.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 static java.util.Collections.singletonList;
  import static org.apache.kafka.clients.consumer.internals.ClassicKafkaConsumer.DEFAULT_REASON;
  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 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 = 3;
  
          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,  MockConsumerInterceptor.class.getName() + ", "
                      + MockConsumerInterceptor.class.getName() + ", "
                      + MockConsumerInterceptor.class.getName());
  
              MockConsumerInterceptor.setThrowOnConfigExceptionThreshold(targetInterceptor);
  
              assertThrows(KafkaException.class, () -> newConsumer(
                      props, new StringDeserializer(), new StringDeserializer()));
  
              assertEquals(3, MockConsumerInterceptor.CONFIG_COUNT.get());
              assertEquals(3, MockConsumerInterceptor.CLOSE_COUNT.get());
  
          } finally {
              MockConsumerInterceptor.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) {
          try {
              newConsumer(groupProtocol, null, Optional.of(Boolean.TRUE));
              fail("Expected an InvalidConfigurationException");
          } catch (InvalidConfigurationException e) {
              // OK, expected
          }
  
          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 e) {
                  // 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);
          consumer.poll(Duration.ofMillis(0));
          // requests: list-offset, fetch
          TestUtils.waitForCondition(() -> {
              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());
          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();
          }
      }
  
      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);
          }
      }
  }
  