/*
 * 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.internals;

import org.apache.kafka.clients.Metadata.LeaderAndEpoch;
import org.apache.kafka.clients.MockClient;
import org.apache.kafka.clients.NodeApiVersions;
import org.apache.kafka.clients.consumer.CloseOptions;
import org.apache.kafka.clients.consumer.ConsumerConfig;
import org.apache.kafka.clients.consumer.ConsumerGroupMetadata;
import org.apache.kafka.clients.consumer.ConsumerRebalanceListener;
import org.apache.kafka.clients.consumer.ConsumerRecord;
import org.apache.kafka.clients.consumer.ConsumerRecords;
import org.apache.kafka.clients.consumer.GroupProtocol;
import org.apache.kafka.clients.consumer.KafkaConsumer;
import org.apache.kafka.clients.consumer.OffsetAndMetadata;
import org.apache.kafka.clients.consumer.OffsetAndTimestamp;
import org.apache.kafka.clients.consumer.OffsetCommitCallback;
import org.apache.kafka.clients.consumer.SubscriptionPattern;
import org.apache.kafka.clients.consumer.internals.events.ApplicationEvent;
import org.apache.kafka.clients.consumer.internals.events.ApplicationEventHandler;
import org.apache.kafka.clients.consumer.internals.events.ApplyAssignmentEvent;
import org.apache.kafka.clients.consumer.internals.events.AssignmentChangeEvent;
import org.apache.kafka.clients.consumer.internals.events.AsyncCommitEvent;
import org.apache.kafka.clients.consumer.internals.events.AsyncPollEvent;
import org.apache.kafka.clients.consumer.internals.events.BackgroundEvent;
import org.apache.kafka.clients.consumer.internals.events.CommitEvent;
import org.apache.kafka.clients.consumer.internals.events.CommitOnCloseEvent;
import org.apache.kafka.clients.consumer.internals.events.CompletableApplicationEvent;
import org.apache.kafka.clients.consumer.internals.events.CompletableBackgroundEvent;
import org.apache.kafka.clients.consumer.internals.events.CompletableEventReaper;
import org.apache.kafka.clients.consumer.internals.events.ConsumerRebalanceListenerCallbackCompletedEvent;
import org.apache.kafka.clients.consumer.internals.events.ErrorEvent;
import org.apache.kafka.clients.consumer.internals.events.EventProcessor;
import org.apache.kafka.clients.consumer.internals.events.FetchCommittedOffsetsEvent;
import org.apache.kafka.clients.consumer.internals.events.LeaveGroupOnCloseEvent;
import org.apache.kafka.clients.consumer.internals.events.ListOffsetsEvent;
import org.apache.kafka.clients.consumer.internals.events.PartitionsAssignedEvent;
import org.apache.kafka.clients.consumer.internals.events.PartitionsRemovedEvent;
import org.apache.kafka.clients.consumer.internals.events.ResetOffsetEvent;
import org.apache.kafka.clients.consumer.internals.events.SeekUnvalidatedEvent;
import org.apache.kafka.clients.consumer.internals.events.StreamsOnTasksAssignedCallbackCompletedEvent;
import org.apache.kafka.clients.consumer.internals.events.StreamsTasksAssignedEvent;
import org.apache.kafka.clients.consumer.internals.events.SyncCommitEvent;
import org.apache.kafka.clients.consumer.internals.events.TopicPatternSubscriptionChangeEvent;
import org.apache.kafka.clients.consumer.internals.events.TopicRe2JPatternSubscriptionChangeEvent;
import org.apache.kafka.clients.consumer.internals.events.TopicSubscriptionChangeEvent;
import org.apache.kafka.clients.consumer.internals.events.UnsubscribeEvent;
import org.apache.kafka.clients.consumer.internals.metrics.AsyncConsumerMetrics;
import org.apache.kafka.common.KafkaException;
import org.apache.kafka.common.Metric;
import org.apache.kafka.common.Node;
import org.apache.kafka.common.TopicPartition;
import org.apache.kafka.common.Uuid;
import org.apache.kafka.common.config.ConfigException;
import org.apache.kafka.common.errors.GroupAuthorizationException;
import org.apache.kafka.common.errors.InterruptException;
import org.apache.kafka.common.errors.InvalidGroupIdException;
import org.apache.kafka.common.errors.TimeoutException;
import org.apache.kafka.common.errors.UnsupportedVersionException;
import org.apache.kafka.common.errors.WakeupException;
import org.apache.kafka.common.internals.ClusterResourceListeners;
import org.apache.kafka.common.message.ConsumerGroupHeartbeatResponseData;
import org.apache.kafka.common.metrics.Metrics;
import org.apache.kafka.common.protocol.ApiKeys;
import org.apache.kafka.common.protocol.Errors;
import org.apache.kafka.common.requests.ConsumerGroupHeartbeatResponse;
import org.apache.kafka.common.requests.FindCoordinatorResponse;
import org.apache.kafka.common.requests.JoinGroupRequest;
import org.apache.kafka.common.requests.ListOffsetsRequest;
import org.apache.kafka.common.requests.MetadataResponse;
import org.apache.kafka.common.requests.RequestTestUtils;
import org.apache.kafka.common.serialization.StringDeserializer;
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.Timer;
import org.apache.kafka.test.MockConsumerInterceptor;
import org.apache.kafka.test.TestUtils;

import org.junit.jupiter.api.AfterEach;
import org.junit.jupiter.api.Test;
import org.junit.jupiter.api.function.Executable;
import org.junit.jupiter.params.ParameterizedTest;
import org.junit.jupiter.params.provider.Arguments;
import org.junit.jupiter.params.provider.MethodSource;
import org.junit.jupiter.params.provider.ValueSource;
import org.mockito.ArgumentCaptor;
import org.mockito.ArgumentMatchers;
import org.mockito.MockedStatic;
import org.mockito.Mockito;

import java.time.Duration;
import java.util.Arrays;
import java.util.Collection;
import java.util.Collections;
import java.util.HashMap;
import java.util.HashSet;
import java.util.List;
import java.util.Locale;
import java.util.Map;
import java.util.Optional;
import java.util.Properties;
import java.util.Set;
import java.util.SortedSet;
import java.util.TreeSet;
import java.util.UUID;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.Future;
import java.util.concurrent.LinkedBlockingQueue;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.concurrent.atomic.AtomicReference;
import java.util.function.Predicate;
import java.util.regex.Pattern;
import java.util.stream.Stream;

import static java.util.Arrays.asList;
import static java.util.Collections.singleton;
import static java.util.Collections.singletonList;
import static org.apache.kafka.clients.consumer.internals.AbstractMembershipManager.TOPIC_PARTITION_COMPARATOR;
import static org.apache.kafka.clients.consumer.internals.ConsumerRebalanceListenerMethodName.ON_PARTITIONS_ASSIGNED;
import static org.apache.kafka.clients.consumer.internals.ConsumerRebalanceListenerMethodName.ON_PARTITIONS_LOST;
import static org.apache.kafka.clients.consumer.internals.ConsumerRebalanceListenerMethodName.ON_PARTITIONS_REVOKED;
import static org.apache.kafka.clients.consumer.internals.ConsumerUtils.CONSUMER_METRIC_GROUP;
import static org.apache.kafka.clients.consumer.internals.ConsumerUtils.THROW_ON_FETCH_STABLE_OFFSET_UNSUPPORTED;
import static org.apache.kafka.test.TestUtils.requiredConsumerConfig;
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.assertNull;
import static org.junit.jupiter.api.Assertions.assertSame;
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.Mockito.atLeast;
import static org.mockito.Mockito.atLeastOnce;
import static org.mockito.Mockito.clearInvocations;
import static org.mockito.Mockito.doAnswer;
import static org.mockito.Mockito.doReturn;
import static org.mockito.Mockito.doThrow;
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.times;
import static org.mockito.Mockito.verify;
import static org.mockito.Mockito.when;

@SuppressWarnings("unchecked")
public class AsyncKafkaConsumerTest {

    private AsyncKafkaConsumer<String, String> consumer = null;
    private Time time = new MockTime(0);
    private final Metrics metrics = new Metrics();
    private final FetchCollector<String, String> fetchCollector = mock(FetchCollector.class);
    private final ApplicationEventHandler applicationEventHandler = mock(ApplicationEventHandler.class);
    private final ConsumerMetadata metadata = mock(ConsumerMetadata.class);
    private final LinkedBlockingQueue<BackgroundEvent> backgroundEventQueue = new LinkedBlockingQueue<>();
    private final CompletableEventReaper backgroundEventReaper = mock(CompletableEventReaper.class);

    @AfterEach
    public void resetAll() {
        backgroundEventQueue.clear();
        if (consumer != null) {
            try {
                consumer.close(CloseOptions.timeout(Duration.ZERO));
            } catch (Exception swallow) {
                // best effort to clean up after each test, but may throw (ex. if callbacks were
                // throwing errors)
            }
        }
        consumer = null;

        Mockito.framework().clearInlineMocks();
        MockConsumerInterceptor.resetCounters();
    }

    private AsyncKafkaConsumer<String, String> newConsumer() {
        final Properties props = requiredConsumerConfig();
        props.put(ConsumerConfig.GROUP_ID_CONFIG, "group-id");
        return newConsumer(props);
    }

    private AsyncKafkaConsumer<String, String> newConsumerWithoutGroupId() {
        final Properties props = requiredConsumerConfig();
        return newConsumer(props);
    }

    private AsyncKafkaConsumer<String, String> newConsumer(Properties props) {
        return newConsumerWithStreamRebalanceData(props, null);
    }

    private AsyncKafkaConsumer<String, String> newConsumerWithStreamRebalanceData(
        Properties props,
        StreamsRebalanceData streamsRebalanceData
    ) {
        // disable auto-commit by default, so we don't need to handle SyncCommitEvent for each case
        if (!props.containsKey(ConsumerConfig.ENABLE_AUTO_COMMIT_CONFIG)) {
            props.put(ConsumerConfig.ENABLE_AUTO_COMMIT_CONFIG, false);
        }
        final ConsumerConfig config = new ConsumerConfig(props);
        return new AsyncKafkaConsumer<>(
            config,
            new StringDeserializer(),
            new StringDeserializer(),
            time,
            (logContext, time, initializationTimeoutMs, applicationEventBlockingQueue, completableEventReaper, applicationEventProcessorSupplier, networkClientDelegateSupplier, requestManagersSupplier, asyncConsumerMetrics) -> applicationEventHandler,
            logContext -> backgroundEventReaper,
            (logContext, consumerMetadata, subscriptionState, fetchConfig, deserializers, fetchMetricsManager, time) -> fetchCollector,
            (consumerConfig, subscriptionState, logContext, clusterResourceListeners) -> metadata,
            backgroundEventQueue,
            Optional.ofNullable(streamsRebalanceData)
        );
    }

    private AsyncKafkaConsumer<String, String> newConsumer(ConsumerConfig config) {
        return new AsyncKafkaConsumer<>(
            config,
            new StringDeserializer(),
            new StringDeserializer(),
            time,
            (logContext, time, initializationTimeoutMs, applicationEventBlockingQueue, completableEventReaper, applicationEventProcessorSupplier, networkClientDelegateSupplier, requestManagersSupplier, asyncConsumerMetrics) -> applicationEventHandler,
            logContext -> backgroundEventReaper,
            (logContext, consumerMetadata, subscriptionState, fetchConfig, deserializers, fetchMetricsManager, time) -> fetchCollector,
            (consumerConfig, subscriptionState, logContext, clusterResourceListeners) -> metadata,
            backgroundEventQueue,
            Optional.empty()
        );
    }

    private AsyncKafkaConsumer<String, String> newConsumer(
        FetchBuffer fetchBuffer,
        ConsumerInterceptors<String, String> interceptors,
        ConsumerRebalanceListenerInvoker rebalanceListenerInvoker,
        SubscriptionState subscriptions) {
        long retryBackoffMs = 100L;
        int requestTimeoutMs = 30000;
        int defaultApiTimeoutMs = 1000;
        return new AsyncKafkaConsumer<>(
            new LogContext(),
            "client-id",
            new Deserializers<>(new StringDeserializer(), new StringDeserializer(), metrics),
            fetchBuffer,
            fetchCollector,
            interceptors,
            time,
            applicationEventHandler,
            backgroundEventQueue,
            backgroundEventReaper,
            rebalanceListenerInvoker,
            metrics,
            subscriptions,
            metadata,
            retryBackoffMs,
            requestTimeoutMs,
            defaultApiTimeoutMs,
            "group-id",
            false,
            new PositionsValidator(new LogContext(), time, subscriptions, metadata));
    }

    @Test
    public void testSuccessfulStartupShutdown() {
        consumer = newConsumer();
        completeUnsubscribeApplicationEventSuccessfully();
        assertDoesNotThrow(() -> consumer.close());
    }

    @Test
    public void testFailOnClosedConsumer() {
        consumer = newConsumer();
        completeUnsubscribeApplicationEventSuccessfully();
        consumer.close();
        final IllegalStateException res = assertThrows(IllegalStateException.class, consumer::assignment);
        assertEquals("This consumer has already been closed.", res.getMessage());
    }

    @Test
    public void testCommitAsyncWithNullCallback() {
        consumer = newConsumer();
        final TopicPartition t0 = new TopicPartition("t0", 2);
        final TopicPartition t1 = new TopicPartition("t0", 3);
        HashMap<TopicPartition, OffsetAndMetadata> offsets = new HashMap<>();
        offsets.put(t0, new OffsetAndMetadata(10L));
        offsets.put(t1, new OffsetAndMetadata(20L));

        markOffsetsReadyForCommitEvent();
        consumer.commitAsync(offsets, null);

        final ArgumentCaptor<AsyncCommitEvent> commitEventCaptor = ArgumentCaptor.forClass(AsyncCommitEvent.class);
        verify(applicationEventHandler).add(commitEventCaptor.capture());
        final AsyncCommitEvent commitEvent = commitEventCaptor.getValue();
        assertTrue(commitEvent.offsets().isPresent());
        assertEquals(offsets, commitEvent.offsets().get());

        commitEvent.future().complete(offsets);
        assertDoesNotThrow(() -> consumer.commitAsync(offsets, null));

        // Clean-up. Close the consumer here as we know it will cause a TimeoutException to be thrown.
        // If we get an error *other* than the TimeoutException, we'll fail the test.
        try {
            Exception e = assertThrows(KafkaException.class, () -> consumer.close(CloseOptions.timeout(Duration.ZERO)));
            assertInstanceOf(TimeoutException.class, e.getCause());
        } finally {
            consumer = null;
        }
    }

    @Test
    public void testCommitAsyncUserSuppliedCallbackNoException() {
        consumer = newConsumer();

        Map<TopicPartition, OffsetAndMetadata> offsets = new HashMap<>();
        offsets.put(new TopicPartition("my-topic", 1), new OffsetAndMetadata(200L));
        completeCommitAsyncApplicationEventSuccessfully();

        MockCommitCallback callback = new MockCommitCallback();
        assertDoesNotThrow(() -> consumer.commitAsync(offsets, callback));
        forceCommitCallbackInvocation();

        assertEquals(1, callback.invoked);
        assertNull(callback.exception);
    }

    @ParameterizedTest
    @MethodSource("commitExceptionSupplier")
    public void testCommitAsyncUserSuppliedCallbackWithException(Exception exception) {
        consumer = newConsumer();

        Map<TopicPartition, OffsetAndMetadata> offsets = new HashMap<>();
        offsets.put(new TopicPartition("my-topic", 1), new OffsetAndMetadata(200L));
        completeCommitAsyncApplicationEventExceptionally(exception);

        MockCommitCallback callback = new MockCommitCallback();
        assertDoesNotThrow(() -> consumer.commitAsync(offsets, callback));
        forceCommitCallbackInvocation();

        assertSame(exception.getClass(), callback.exception.getClass());
    }

    @Test
    public void testCommitAsyncShouldCopyOffsets() {
        consumer = newConsumer();

        TopicPartition tp = new TopicPartition("t0", 2);
        Map<TopicPartition, OffsetAndMetadata> offsets = new HashMap<>();
        offsets.put(tp, new OffsetAndMetadata(10L));

        markOffsetsReadyForCommitEvent();
        consumer.commitAsync(offsets, null);

        final ArgumentCaptor<AsyncCommitEvent> commitEventCaptor = ArgumentCaptor.forClass(AsyncCommitEvent.class);
        verify(applicationEventHandler).add(commitEventCaptor.capture());
        final AsyncCommitEvent commitEvent = commitEventCaptor.getValue();
        assertTrue(commitEvent.offsets().isPresent());
        assertTrue(commitEvent.offsets().get().containsKey(tp));
        offsets.remove(tp);
        assertTrue(commitEvent.offsets().get().containsKey(tp));
    }

    private static Stream<Exception> commitExceptionSupplier() {
        return Stream.of(
                new KafkaException("Test exception"),
                new GroupAuthorizationException("Group authorization exception"));
    }

    @Test
    public void testCommitted() {
        time = new MockTime(1);
        consumer = newConsumer();
        Map<TopicPartition, OffsetAndMetadata> topicPartitionOffsets = mockTopicPartitionOffset();
        completeFetchedCommittedOffsetApplicationEventSuccessfully(topicPartitionOffsets);

        assertEquals(topicPartitionOffsets, consumer.committed(topicPartitionOffsets.keySet(), Duration.ofMillis(1000)));
        verify(applicationEventHandler).addAndGet(ArgumentMatchers.isA(FetchCommittedOffsetsEvent.class));
        final Metric metric = consumer.metrics()
            .get(consumer.metricsRegistry().metricName("committed-time-ns-total", CONSUMER_METRIC_GROUP));
        assertTrue((double) metric.metricValue() > 0);
    }

    @Test
    public void testCommittedExceptionThrown() {
        consumer = newConsumer();
        Map<TopicPartition, OffsetAndMetadata> offsets = mockTopicPartitionOffset();
        when(applicationEventHandler.addAndGet(
            any(FetchCommittedOffsetsEvent.class))).thenAnswer(invocation -> {
                CompletableApplicationEvent<?> event = invocation.getArgument(0);
                assertInstanceOf(FetchCommittedOffsetsEvent.class, event);
                throw new KafkaException("Test exception");
            });

        assertThrows(KafkaException.class, () -> consumer.committed(offsets.keySet(), Duration.ofMillis(1000)));
    }

    @Test
    public void testWakeupBeforeCallingPoll() {
        consumer = newConsumer();
        final String topicName = "foo";
        final int partition = 3;
        final TopicPartition tp = new TopicPartition(topicName, partition);
        doReturn(Fetch.empty()).when(fetchCollector).collectFetch(any(FetchBuffer.class));
        doReturn(LeaderAndEpoch.noLeaderOrEpoch()).when(metadata).currentLeader(any());
        completeAssignmentChangeEventSuccessfully();
        consumer.assign(singleton(tp));

        consumer.wakeup();

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

    @Test
    public void testWakeupAfterEmptyFetch() {
        consumer = newConsumer();
        final String topicName = "foo";
        final int partition = 3;
        final TopicPartition tp = new TopicPartition(topicName, partition);
        doAnswer(invocation -> {
            consumer.wakeup();
            return Fetch.empty();
        }).doAnswer(invocation -> Fetch.empty()).when(fetchCollector).collectFetch(any(FetchBuffer.class));
        doReturn(LeaderAndEpoch.noLeaderOrEpoch()).when(metadata).currentLeader(any());
        completeAssignmentChangeEventSuccessfully();
        consumer.assign(singleton(tp));

        completeAsyncPollEventSuccessfully();
        assertThrows(WakeupException.class, () -> consumer.poll(Duration.ofMinutes(1)));
        assertDoesNotThrow(() -> consumer.poll(Duration.ZERO));
    }

    @Test
    public void testWakeupAfterNonEmptyFetch() {
        consumer = newConsumer();
        final String topicName = "foo";
        final int partition = 3;
        final TopicPartition tp = new TopicPartition(topicName, partition);
        final List<ConsumerRecord<String, String>> records = asList(
            new ConsumerRecord<>(topicName, partition, 2, "key1", "value1"),
            new ConsumerRecord<>(topicName, partition, 3, "key2", "value2")
        );
        doAnswer(invocation -> {
            consumer.wakeup();
            return Fetch.forPartition(tp, records, true, new OffsetAndMetadata(4, Optional.of(0), ""));
        }).when(fetchCollector).collectFetch(Mockito.any(FetchBuffer.class));
        doReturn(LeaderAndEpoch.noLeaderOrEpoch()).when(metadata).currentLeader(any());
        completeAssignmentChangeEventSuccessfully();
        consumer.assign(singleton(tp));

        completeAsyncPollEventSuccessfully();
        // since wakeup() is called when the non-empty fetch is returned the wakeup should be ignored
        assertDoesNotThrow(() -> consumer.poll(Duration.ofMinutes(1)));
        // the previously ignored wake-up should not be ignored in the next call
        assertThrows(WakeupException.class, () -> consumer.poll(Duration.ZERO));
    }

    @Test
    public void testWakeupWhileWaitingOnReconciliationCheck() {
        FetchBuffer fetchBuffer = mock(FetchBuffer.class);
        SubscriptionState subscriptions = new SubscriptionState(new LogContext(), AutoOffsetResetStrategy.NONE);
        consumer = newConsumer(fetchBuffer, mock(ConsumerInterceptors.class),
            mock(ConsumerRebalanceListenerInvoker.class), subscriptions);

        final TopicPartition tp = new TopicPartition("topic1", 0);
        subscriptions.assignFromUser(singleton(tp));
        subscriptions.seek(tp, 0);

        // Do not complete the AsyncPollEvent and call wakeup().
        // The call to poll should throw WakeupException without blocking for the full timeout.
        doAnswer(invocation -> {
            consumer.wakeup();
            return null;
        }).when(applicationEventHandler).add(ArgumentMatchers.isA(AsyncPollEvent.class));
        doReturn(Fetch.empty()).when(fetchCollector).collectFetch(any(FetchBuffer.class));

        long startTime = System.currentTimeMillis();
        assertThrows(WakeupException.class, () -> consumer.poll(Duration.ofMillis(1500)));
        long elapsed = System.currentTimeMillis() - startTime;

        assertTrue(elapsed < 500, "Wakeup should interrupt promptly, took " + elapsed + "ms");
    }

    @Test
    public void testCommitInRebalanceCallback() {
        consumer = newConsumer();
        final String topicName = "foo";
        final int partition = 3;
        final TopicPartition tp = new TopicPartition(topicName, partition);
        doAnswer(invocation -> Fetch.empty()).when(fetchCollector).collectFetch(Mockito.any(FetchBuffer.class));
        SortedSet<TopicPartition> sortedPartitions = new TreeSet<>(TOPIC_PARTITION_COMPARATOR);
        sortedPartitions.add(tp);
        CompletableBackgroundEvent<Void> e = new PartitionsRemovedEvent(ON_PARTITIONS_REVOKED, sortedPartitions);
        backgroundEventQueue.add(e);
        completeCommitSyncApplicationEventSuccessfully();
        final AtomicBoolean callbackExecuted = new AtomicBoolean(false);

        ConsumerRebalanceListener listener = new ConsumerRebalanceListener() {
            @Override
            public void onPartitionsRevoked(final Collection<TopicPartition> partitions) {
                assertDoesNotThrow(() -> consumer.commitSync(Map.of(tp, new OffsetAndMetadata(0))));
                callbackExecuted.set(true);
            }

            @Override
            public void onPartitionsAssigned(final Collection<TopicPartition> partitions) {
                // no-op
            }
        };

        completeTopicSubscriptionChangeEventSuccessfully();
        consumer.subscribe(Collections.singletonList(topicName), listener);
        completeAsyncPollEventSuccessfully();
        consumer.poll(Duration.ZERO);
        assertTrue(callbackExecuted.get());
    }

    @Test
    public void testClearWakeupTriggerAfterPoll() {
        consumer = newConsumer();
        final String topicName = "foo";
        final int partition = 3;
        final TopicPartition tp = new TopicPartition(topicName, partition);
        final List<ConsumerRecord<String, String>> records = asList(
            new ConsumerRecord<>(topicName, partition, 2, "key1", "value1"),
            new ConsumerRecord<>(topicName, partition, 3, "key2", "value2")
        );
        doReturn(Fetch.forPartition(tp, records, true, new OffsetAndMetadata(4, Optional.of(0), "")))
            .when(fetchCollector).collectFetch(any(FetchBuffer.class));
        doReturn(LeaderAndEpoch.noLeaderOrEpoch()).when(metadata).currentLeader(any());
        completeAssignmentChangeEventSuccessfully();
        consumer.assign(singleton(tp));

        completeAsyncPollEventSuccessfully();
        consumer.poll(Duration.ZERO);

        assertDoesNotThrow(() -> consumer.poll(Duration.ZERO));
    }

    /**
     * Test that poll() does not return records until the reconciliation check is complete.
     * This prevents a race condition where records could be returned for partitions that
     * are being revoked (see KAFKA-20332).
     */
    @Test
    public void testPollWaitsForReconciliationCheckComplete() {
        final String topicName = "foo";
        final int partition = 3;
        final TopicPartition tp = new TopicPartition(topicName, partition);
        final List<ConsumerRecord<String, String>> records = asList(
            new ConsumerRecord<>(topicName, partition, 2, "key1", "value1"),
            new ConsumerRecord<>(topicName, partition, 3, "key2", "value2")
        );

        SubscriptionState subscriptions = new SubscriptionState(new LogContext(), AutoOffsetResetStrategy.NONE);
        consumer = newConsumer(
            mock(FetchBuffer.class),
            new ConsumerInterceptors<>(Collections.emptyList(), metrics),
            mock(ConsumerRebalanceListenerInvoker.class),
            subscriptions);

        doReturn(LeaderAndEpoch.noLeaderOrEpoch()).when(metadata).currentLeader(any());
        // PositionsValidator starts with metadataUpdateVersion=-1. Stub metadata.updateVersion() to match,
        // so canSkipUpdateFetchPositions() passes and we test the reconciliation check path.
        doReturn(-1).when(metadata).updateVersion();

        completeTopicSubscriptionChangeEventSuccessfully();
        consumer.subscribe(singleton(topicName), mock(ConsumerRebalanceListener.class));
        // Simulate partition assignment from group coordinator
        subscriptions.assignFromSubscribed(singleton(tp));

        // Set up position so canSkipUpdateFetchPositions() returns true (partition in FETCHING state)
        completeSeekUnvalidatedEventSuccessfully();
        subscriptions.seek(tp, 0);

        // Set up fetch collector to return records when called
        doReturn(Fetch.forPartition(tp, records, true, new OffsetAndMetadata(4, Optional.of(0), "")))
            .when(fetchCollector).collectFetch(any(FetchBuffer.class));

        // Capture the AsyncPollEvent to manually control when reconciliation check is marked complete
        AtomicReference<AsyncPollEvent> capturedEvent = new AtomicReference<>();
        doAnswer(invocation -> {
            AsyncPollEvent event = invocation.getArgument(0);
            assertTrue(capturedEvent.compareAndSet(null, event));
            // Do NOT mark reconciliation check complete - simulating background hasn't processed it yet
            return null;
        }).when(applicationEventHandler).add(ArgumentMatchers.isA(AsyncPollEvent.class));
        consumer.setHasPendingReconciliation(true);

        // Poll should return empty because reconciliation check is not complete.
        ConsumerRecords<?, ?> result1 = consumer.poll(Duration.ZERO);
        assertTrue(result1.isEmpty(), "Poll should not return records if it hasn't completed checking and triggering pending reconciliations.");

        // Now mark reconciliation check complete on the captured event
        assertNotNull(capturedEvent.get(), "AsyncPollEvent should have been captured");
        capturedEvent.get().markReconciliationCheckComplete();

        // Next poll should return the records since reconciliation check is now complete
        ConsumerRecords<?, ?> result2 = consumer.poll(Duration.ZERO);
        assertEquals(2, result2.count(), "Expected 2 records after reconciliation check is complete");
    }

    @Test
    public void testPollDoesNotWaitForReconciliationCheckIfNoPendingReconciliation() {
        final String topicName = "foo";
        final int partition = 3;
        final TopicPartition tp = new TopicPartition(topicName, partition);
        final List<ConsumerRecord<String, String>> records = asList(
                new ConsumerRecord<>(topicName, partition, 2, "key1", "value1"),
                new ConsumerRecord<>(topicName, partition, 3, "key2", "value2")
        );

        SubscriptionState subscriptions = new SubscriptionState(new LogContext(), AutoOffsetResetStrategy.NONE);
        consumer = newConsumer(
                mock(FetchBuffer.class),
                new ConsumerInterceptors<>(Collections.emptyList(), metrics),
                mock(ConsumerRebalanceListenerInvoker.class),
                subscriptions);

        doReturn(LeaderAndEpoch.noLeaderOrEpoch()).when(metadata).currentLeader(any());
        // PositionsValidator starts with metadataUpdateVersion=-1. Stub metadata.updateVersion() to match,
        // so canSkipUpdateFetchPositions() passes and we test the reconciliation check path.
        doReturn(-1).when(metadata).updateVersion();

        completeTopicSubscriptionChangeEventSuccessfully();
        consumer.subscribe(singleton(topicName), mock(ConsumerRebalanceListener.class));
        // Simulate partition assignment from group coordinator
        subscriptions.assignFromSubscribed(singleton(tp));

        // Set up position so canSkipUpdateFetchPositions() returns true (partition in FETCHING state)
        completeSeekUnvalidatedEventSuccessfully();
        subscriptions.seek(tp, 0);

        // Set up fetch collector to return records when called
        doReturn(Fetch.forPartition(tp, records, true, new OffsetAndMetadata(4, Optional.of(0), "")))
                .when(fetchCollector).collectFetch(any(FetchBuffer.class));

        // Capture the AsyncPollEvent but leave the reconciliation check incomplete.
        // Since there is no pending reconciliation, poll should not wait for it.
        AtomicReference<AsyncPollEvent> capturedEvent = new AtomicReference<>();
        doAnswer(invocation -> {
            AsyncPollEvent event = invocation.getArgument(0);
            assertTrue(capturedEvent.compareAndSet(null, event));
            // Do NOT mark reconciliation check complete - simulating background hasn't processed it yet
            return null;
        }).when(applicationEventHandler).add(ArgumentMatchers.isA(AsyncPollEvent.class));
        consumer.setHasPendingReconciliation(false);
        
        // Poll does not wait AsyncPollEvent if there is no pending reconciliation.
        ConsumerRecords<?, ?> result = consumer.poll(Duration.ZERO);

        assertNotNull(capturedEvent.get(), "AsyncPollEvent should have been captured");
        assertFalse(capturedEvent.get().isReconciliationCheckComplete(), "Reconciliation check should still be incomplete");
        assertEquals(2, result.count(), "Expected records without waiting when no reconciliation is pending");
    }

    @Test
    public void testEnsureCallbackExecutedByApplicationThread() {
        consumer = newConsumer();
        final String currentThread = Thread.currentThread().getName();
        MockCommitCallback callback = new MockCommitCallback();
        completeCommitAsyncApplicationEventSuccessfully();

        assertDoesNotThrow(() -> consumer.commitAsync(new HashMap<>(), callback));
        forceCommitCallbackInvocation();
        assertEquals(1, callback.invoked);
        assertEquals(currentThread, callback.completionThread);
    }

    @Test
    public void testEnsureCommitSyncExecutedCommitAsyncCallbacks() {
        consumer = newConsumer();
        KafkaException callbackException = new KafkaException("Async commit callback failed");
        OffsetCommitCallback callback = (offsets, exception) -> {
            throw callbackException;
        };

        assertDoesNotThrow(() -> consumer.commitAsync(new HashMap<>(), callback));
        assertThrows(callbackException.getClass(), () -> consumer.commitSync());
    }

    @Test
    public void testCommitSyncAwaitsCommitAsyncCompletionWithEmptyOffsets() {
        final TopicPartition tp = new TopicPartition("foo", 0);
        final CompletableFuture<Void> asyncCommitFuture = setUpConsumerWithIncompleteAsyncCommit(tp);

        // Commit async is not completed yet, so commit sync should wait for it to complete (time out)
        assertThrows(TimeoutException.class, () -> consumer.commitSync(Collections.emptyMap(), Duration.ofMillis(100)));

        // Complete exceptionally async commit event
        asyncCommitFuture.completeExceptionally(new KafkaException("Test exception"));

        // Commit async is completed, so commit sync completes immediately (since offsets are empty)
        assertDoesNotThrow(() -> consumer.commitSync(Collections.emptyMap(), Duration.ofMillis(100)));
    }

    @Test
    public void testCommitSyncAwaitsCommitAsyncCompletionWithNonEmptyOffsets() {
        final TopicPartition tp = new TopicPartition("foo", 0);
        final CompletableFuture<Void> asyncCommitFuture = setUpConsumerWithIncompleteAsyncCommit(tp);

        // Mock to complete sync event
        completeCommitSyncApplicationEventSuccessfully();

        // Commit async is not completed yet, so commit sync should wait for it to complete (time out)
        assertThrows(TimeoutException.class, () -> consumer.commitSync(Collections.singletonMap(tp, new OffsetAndMetadata(20)), Duration.ofMillis(100)));

        // Complete async commit event
        asyncCommitFuture.complete(null);

        // Commit async is completed, so commit sync does not need to wait before committing its offsets
        assertDoesNotThrow(() -> consumer.commitSync(Collections.singletonMap(tp, new OffsetAndMetadata(20)), Duration.ofMillis(100)));
    }

    @Test
    public void testCommitSyncAwaitsCommitAsyncButDoesNotFail() {
        final TopicPartition tp = new TopicPartition("foo", 0);
        final CompletableFuture<Void> asyncCommitFuture = setUpConsumerWithIncompleteAsyncCommit(tp);

        // Mock to complete sync event
        completeCommitSyncApplicationEventSuccessfully();

        // Commit async is not completed yet, so commit sync should wait for it to complete (time out)
        assertThrows(TimeoutException.class, () -> consumer.commitSync(Collections.singletonMap(tp, new OffsetAndMetadata(20)), Duration.ofMillis(100)));

        // Complete exceptionally async commit event
        asyncCommitFuture.completeExceptionally(new KafkaException("Test exception"));

        // Commit async is completed exceptionally, but this will be handled by commit callback - commit sync should not fail.
        assertDoesNotThrow(() -> consumer.commitSync(Collections.singletonMap(tp, new OffsetAndMetadata(20)), Duration.ofMillis(100)));
    }

    @Test
    public void testCommitSyncShouldCopyOffsets() {
        consumer = newConsumer();

        TopicPartition tp = new TopicPartition("t0", 2);
        Map<TopicPartition, OffsetAndMetadata> offsets = new HashMap<>();
        offsets.put(tp, new OffsetAndMetadata(10L));

        completeCommitSyncApplicationEventSuccessfully();
        consumer.commitSync(offsets);

        final ArgumentCaptor<SyncCommitEvent> commitEventCaptor = ArgumentCaptor.forClass(SyncCommitEvent.class);
        verify(applicationEventHandler).add(commitEventCaptor.capture());
        final SyncCommitEvent commitEvent = commitEventCaptor.getValue();
        assertTrue(commitEvent.offsets().isPresent());
        assertTrue(commitEvent.offsets().get().containsKey(tp));
        offsets.remove(tp);
        assertTrue(commitEvent.offsets().get().containsKey(tp));
    }

    private CompletableFuture<Void> setUpConsumerWithIncompleteAsyncCommit(TopicPartition tp) {
        time = new MockTime(1);
        consumer = newConsumer();

        // Commit async (incomplete)
        doReturn(LeaderAndEpoch.noLeaderOrEpoch()).when(metadata).currentLeader(any());
        completeAssignmentChangeEventSuccessfully();
        consumer.assign(Collections.singleton(tp));
        completeSeekUnvalidatedEventSuccessfully();
        consumer.seek(tp, 20);
        markOffsetsReadyForCommitEvent();
        consumer.commitAsync();

        CompletableApplicationEvent<Void> event = getLastEnqueuedEvent();
        return event.future();
    }

    // ArgumentCaptor's type-matching does not work reliably with Java 8, so we cannot directly capture the AsyncCommitEvent
    // Instead, we capture the super-class CompletableApplicationEvent and fetch the last captured event.
    private <T> CompletableApplicationEvent<T> getLastEnqueuedEvent() {
        final ArgumentCaptor<CompletableApplicationEvent<T>> eventArgumentCaptor = ArgumentCaptor.forClass(CompletableApplicationEvent.class);
        verify(applicationEventHandler, atLeast(1)).add(eventArgumentCaptor.capture());
        final List<CompletableApplicationEvent<T>> allValues = eventArgumentCaptor.getAllValues();
        return allValues.get(allValues.size() - 1);
    }

    private <T> CompletableApplicationEvent<T> addAndGetLastEnqueuedEvent() {
        final ArgumentCaptor<CompletableApplicationEvent<T>> eventArgumentCaptor = ArgumentCaptor.forClass(CompletableApplicationEvent.class);
        verify(applicationEventHandler, atLeast(1)).addAndGet(eventArgumentCaptor.capture());
        final List<CompletableApplicationEvent<T>> allValues = eventArgumentCaptor.getAllValues();
        return allValues.get(allValues.size() - 1);
    }

    @Test
    public void testEnsurePollExecutedCommitAsyncCallbacks() {
        consumer = newConsumer();
        MockCommitCallback callback = new MockCommitCallback();
        completeCommitAsyncApplicationEventSuccessfully();
        doReturn(Fetch.empty()).when(fetchCollector).collectFetch(any(FetchBuffer.class));

        completeAssignmentChangeEventSuccessfully();
        consumer.assign(Collections.singleton(new TopicPartition("foo", 0)));
        assertDoesNotThrow(() -> consumer.commitAsync(new HashMap<>(), callback));
        completeAsyncPollEventSuccessfully();
        assertMockCommitCallbackInvoked(() -> consumer.poll(Duration.ZERO), callback);
    }

    @Test
    public void testEnsureShutdownExecutedCommitAsyncCallbacks() {
        consumer = newConsumer();
        completeUnsubscribeApplicationEventSuccessfully();
        MockCommitCallback callback = new MockCommitCallback();
        completeCommitAsyncApplicationEventSuccessfully();
        assertDoesNotThrow(() -> consumer.commitAsync(new HashMap<>(), callback));
        assertMockCommitCallbackInvoked(() -> consumer.close(), callback);
    }

    @Test
    public void testVerifyApplicationEventOnShutdown() {
        consumer = newConsumer();
        completeUnsubscribeApplicationEventSuccessfully();
        doReturn(null).when(applicationEventHandler).addAndGet(any());
        consumer.close();
        verify(applicationEventHandler).add(any(CommitOnCloseEvent.class));
        verify(applicationEventHandler).addAndGet(any(LeaveGroupOnCloseEvent.class));
    }

    @ParameterizedTest
    @ValueSource(longs = {0, ConsumerUtils.DEFAULT_CLOSE_TIMEOUT_MS})
    public void testCloseLeavesGroup(long timeoutMs) {
        SubscriptionState subscriptions = mock(SubscriptionState.class);
        consumer = spy(newConsumer(
            mock(FetchBuffer.class),
            mock(ConsumerInterceptors.class),
            mock(ConsumerRebalanceListenerInvoker.class),
            subscriptions));
        consumer.close(CloseOptions.timeout(Duration.ofMillis(timeoutMs)));
        verify(applicationEventHandler).addAndGet(any(LeaveGroupOnCloseEvent.class));
    }

    @Test
    public void testCloseLeavesGroupDespiteOnPartitionsLostError() {
        // If rebalance listener failed to execute during close, we still send the leave group,
        // and proceed with closing the consumer.
        Throwable rootError = new KafkaException("Intentional error");
        Set<TopicPartition> partitions = singleton(new TopicPartition("topic1", 0));
        SubscriptionState subscriptions = mock(SubscriptionState.class);
        when(subscriptions.assignedPartitions()).thenReturn(partitions);
        ConsumerRebalanceListenerInvoker invoker = mock(ConsumerRebalanceListenerInvoker.class);
        doAnswer(invocation -> rootError).when(invoker).invokePartitionsLost(any(SortedSet.class));

        consumer = spy(newConsumer(
            mock(FetchBuffer.class),
            new ConsumerInterceptors<>(Collections.emptyList(), metrics),
            invoker,
            subscriptions));
        consumer.setGroupAssignmentSnapshot(partitions);

        Throwable t = assertThrows(KafkaException.class, () -> consumer.close(CloseOptions.timeout(Duration.ZERO)));
        assertNotNull(t.getCause());
        assertEquals(rootError, t.getCause());

        verify(applicationEventHandler).addAndGet(any(LeaveGroupOnCloseEvent.class));
    }

    @ParameterizedTest
    @ValueSource(longs = {0, ConsumerUtils.DEFAULT_CLOSE_TIMEOUT_MS})
    public void testCloseLeavesGroupDespiteInterrupt(long timeoutMs) {
        Set<TopicPartition> partitions = singleton(new TopicPartition("topic1", 0));
        SubscriptionState subscriptions = mock(SubscriptionState.class);
        when(subscriptions.assignedPartitions()).thenReturn(partitions);
        when(applicationEventHandler.addAndGet(any(CompletableApplicationEvent.class))).thenThrow(InterruptException.class);
        consumer = spy(newConsumer(
            mock(FetchBuffer.class),
            mock(ConsumerInterceptors.class),
            mock(ConsumerRebalanceListenerInvoker.class),
            subscriptions));

        Duration timeout = Duration.ofMillis(timeoutMs);

        try {
            assertThrows(InterruptException.class, () -> consumer.close(CloseOptions.timeout(timeout)));
        } finally {
            Thread.interrupted();
        }

        verify(applicationEventHandler).add(any(CommitOnCloseEvent.class));
        verify(applicationEventHandler).addAndGet(any(LeaveGroupOnCloseEvent.class));
    }

    @Test
    public void testCommitSyncAllConsumed() {
        SubscriptionState subscriptions = new SubscriptionState(new LogContext(), AutoOffsetResetStrategy.NONE);
        consumer = newConsumer(
            mock(FetchBuffer.class),
            mock(ConsumerInterceptors.class),
            mock(ConsumerRebalanceListenerInvoker.class),
            subscriptions);
        completeTopicSubscriptionChangeEventSuccessfully();
        consumer.subscribe(singleton("topic"), mock(ConsumerRebalanceListener.class));
        subscriptions.assignFromSubscribed(singleton(new TopicPartition("topic", 0)));
        completeSeekUnvalidatedEventSuccessfully();
        subscriptions.seek(new TopicPartition("topic", 0), 100);
        markOffsetsReadyForCommitEvent();
        consumer.commitSyncAllConsumed(time.timer(100));

        ArgumentCaptor<SyncCommitEvent> eventCaptor = ArgumentCaptor.forClass(SyncCommitEvent.class);
        verify(applicationEventHandler).add(eventCaptor.capture());
        SyncCommitEvent capturedEvent = eventCaptor.getValue();
        assertFalse(capturedEvent.offsets().isPresent(), "Expected empty optional offsets");
    }

    @Test
    public void testAutoCommitSyncDisabled() {
        SubscriptionState subscriptions = new SubscriptionState(new LogContext(), AutoOffsetResetStrategy.NONE);
        consumer = newConsumer(
            mock(FetchBuffer.class),
            mock(ConsumerInterceptors.class),
            mock(ConsumerRebalanceListenerInvoker.class),
            subscriptions);
        completeTopicSubscriptionChangeEventSuccessfully();
        consumer.subscribe(singleton("topic"), mock(ConsumerRebalanceListener.class));
        subscriptions.assignFromSubscribed(singleton(new TopicPartition("topic", 0)));
        completeSeekUnvalidatedEventSuccessfully();
        subscriptions.seek(new TopicPartition("topic", 0), 100);
        completeUnsubscribeApplicationEventSuccessfully();
        consumer.close();
        verify(applicationEventHandler, never()).add(any(SyncCommitEvent.class));
    }

    private void assertMockCommitCallbackInvoked(final Executable task, final MockCommitCallback callback) {
        assertDoesNotThrow(task);
        assertEquals(1, callback.invoked);
        assertNull(callback.exception);
    }

    private static class MockCommitCallback implements OffsetCommitCallback {
        public int invoked = 0;
        public Exception exception = null;
        public String completionThread;

        @Override
        public void onComplete(Map<TopicPartition, OffsetAndMetadata> offsets, Exception exception) {
            invoked++;
            this.completionThread = Thread.currentThread().getName();
            this.exception = exception;
        }
    }

    @Test
    public void testAssign() {
        consumer = newConsumer();
        final TopicPartition tp = new TopicPartition("foo", 3);
        completeAssignmentChangeEventSuccessfully();
        consumer.assign(singleton(tp));
        assertTrue(consumer.subscription().isEmpty());
        assertTrue(consumer.assignment().contains(tp));
        verify(applicationEventHandler).addAndGet(any(AssignmentChangeEvent.class));
    }

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

    @Test
    public void testAssignOnEmptyTopicPartition() {
        consumer = newConsumer();
        completeUnsubscribeApplicationEventSuccessfully();

        consumer.assign(Collections.emptyList());
        assertTrue(consumer.subscription().isEmpty());
        assertTrue(consumer.assignment().isEmpty());
    }

    @Test
    public void testAssignOnNullTopicInPartition() {
        consumer = newConsumer();
        assertThrows(IllegalArgumentException.class, () -> consumer.assign(singleton(new TopicPartition(null, 0))));
    }

    @Test
    public void testAssignOnEmptyTopicInPartition() {
        consumer = newConsumer();
        assertThrows(IllegalArgumentException.class, () -> consumer.assign(singleton(new TopicPartition("  ", 0))));
    }

    @Test
    public void testBeginningOffsetsFailsIfNullPartitions() {
        consumer = newConsumer();
        assertThrows(NullPointerException.class, () -> consumer.beginningOffsets(null,
            Duration.ofMillis(1)));
    }

    @Test
    public void testBeginningOffsets() {
        consumer = newConsumer();
        Map<TopicPartition, OffsetAndTimestampInternal> expectedOffsets = mockOffsetAndTimestamp();

        when(applicationEventHandler.addAndGet(any(ListOffsetsEvent.class))).thenAnswer(invocation -> {
            ListOffsetsEvent event = invocation.getArgument(0);
            Timer timer = time.timer(event.deadlineMs() - time.milliseconds());
            if (timer.remainingMs() == 0) {
                fail("Timer duration should not be zero.");
            }
            return expectedOffsets;
        });

        Map<TopicPartition, Long> result = assertDoesNotThrow(() -> consumer.beginningOffsets(expectedOffsets.keySet(), Duration.ofMillis(1)));

        expectedOffsets.forEach((key, value) -> {
            assertTrue(result.containsKey(key));
            assertEquals(value.offset(), result.get(key));
        });
        verify(applicationEventHandler).addAndGet(any(ListOffsetsEvent.class));
    }

    @Test
    public void testBeginningOffsetsThrowsKafkaExceptionForUnderlyingExecutionFailure() {
        consumer = newConsumer();
        Set<TopicPartition> partitions = mockTopicPartitionOffset().keySet();
        Throwable eventProcessingFailure = new KafkaException("Unexpected failure " +
            "processing List Offsets event");
        doThrow(eventProcessingFailure).when(applicationEventHandler).addAndGet(
            any(ListOffsetsEvent.class));
        Throwable consumerError = assertThrows(KafkaException.class,
            () -> consumer.beginningOffsets(partitions,
                Duration.ofMillis(1)));
        assertEquals(eventProcessingFailure, consumerError);
        verify(applicationEventHandler).addAndGet(ArgumentMatchers.isA(ListOffsetsEvent.class));
    }

    @Test
    public void testBeginningOffsetsTimeoutOnEventProcessingTimeout() {
        consumer = newConsumer();
        doThrow(new TimeoutException()).when(applicationEventHandler).addAndGet(any());
        assertThrows(TimeoutException.class,
            () -> consumer.beginningOffsets(
                Collections.singletonList(new TopicPartition("t1", 0)),
                Duration.ofMillis(1)));
        verify(applicationEventHandler).addAndGet(ArgumentMatchers.isA(ListOffsetsEvent.class));
    }

    @Test
    public void testOffsetsForTimesOnNullPartitions() {
        consumer = newConsumer();
        assertThrows(NullPointerException.class, () -> consumer.offsetsForTimes(null,
            Duration.ofMillis(1)));
    }

    @Test
    public void testOffsetsForTimesFailsOnNegativeTargetTimes() {
        consumer = newConsumer();
        assertThrows(IllegalArgumentException.class,
                () -> consumer.offsetsForTimes(Collections.singletonMap(new TopicPartition(
                                "topic1", 1), ListOffsetsRequest.EARLIEST_TIMESTAMP),
                        Duration.ofMillis(1)));

        assertThrows(IllegalArgumentException.class,
                () -> consumer.offsetsForTimes(Collections.singletonMap(new TopicPartition(
                                "topic1", 1), ListOffsetsRequest.LATEST_TIMESTAMP),
                        Duration.ofMillis(1)));

        assertThrows(IllegalArgumentException.class,
                () -> consumer.offsetsForTimes(Collections.singletonMap(new TopicPartition(
                                "topic1", 1), ListOffsetsRequest.MAX_TIMESTAMP),
                        Duration.ofMillis(1)));
    }

    @Test
    public void testOffsetsForTimes() {
        consumer = newConsumer();
        Map<TopicPartition, OffsetAndTimestampInternal> expectedResult = mockOffsetAndTimestamp();
        Map<TopicPartition, Long> timestampToSearch = mockTimestampToSearch();

        doReturn(expectedResult).when(applicationEventHandler).addAndGet(any());
        Map<TopicPartition, OffsetAndTimestamp> result =
                assertDoesNotThrow(() -> consumer.offsetsForTimes(timestampToSearch, Duration.ofMillis(1)));
        expectedResult.forEach((key, value) -> {
            OffsetAndTimestamp expected = value.buildOffsetAndTimestamp();
            assertEquals(expected, result.get(key));
        });
        verify(applicationEventHandler).addAndGet(ArgumentMatchers.isA(ListOffsetsEvent.class));
    }

    @Test
    public void testOffsetsForTimesTimeoutException() {
        consumer = newConsumer();
        long timeout = 100;
        doThrow(new TimeoutException("Event did not complete in time and was expired by the reaper"))
            .when(applicationEventHandler).addAndGet(any());

        Throwable t = assertThrows(
            TimeoutException.class,
            () -> consumer.offsetsForTimes(mockTimestampToSearch(), Duration.ofMillis(timeout)));
        assertEquals("Failed to get offsets by times in " + timeout + "ms", t.getMessage());
    }

    @Test
    public void testBeginningOffsetsTimeoutException() {
        consumer = newConsumer();
        long timeout = 100;
        doThrow(new TimeoutException("Event did not complete in time and was expired by the reaper"))
            .when(applicationEventHandler).addAndGet(any());

        Throwable t = assertThrows(
            TimeoutException.class,
            () -> consumer.beginningOffsets(Collections.singleton(new TopicPartition("topic", 5)),
                Duration.ofMillis(timeout)));
        assertEquals("Failed to get offsets by times in " + timeout + "ms", t.getMessage());
    }

    @Test
    public void testEndOffsetsTimeoutException() {
        consumer = newConsumer();
        long timeout = 100;
        doThrow(new TimeoutException("Event did not complete in time and was expired by the reaper"))
            .when(applicationEventHandler).addAndGet(any());

        Throwable t = assertThrows(
            TimeoutException.class,
            () -> consumer.endOffsets(Collections.singleton(new TopicPartition("topic", 5)),
                Duration.ofMillis(timeout)));
        assertEquals("Failed to get offsets by times in " + timeout + "ms", t.getMessage());
    }

    // This test ensures same behaviour as the current consumer when offsetsForTimes is called
    // with 0 timeout. It should return map with all requested partitions as keys, with null
    // OffsetAndTimestamp as value.
    @Test
    public void testBeginningOffsetsWithZeroTimeout() {
        consumer = newConsumer();
        TopicPartition tp = new TopicPartition("topic1", 0);
        Map<TopicPartition, Long> result =
                assertDoesNotThrow(() -> consumer.beginningOffsets(Collections.singletonList(tp), Duration.ZERO));
        assertNotNull(result);
        assertEquals(0, result.size());
        verify(applicationEventHandler).add(ArgumentMatchers.isA(ListOffsetsEvent.class));
    }

    @Test
    public void testOffsetsForTimesWithZeroTimeout() {
        consumer = newConsumer();
        TopicPartition tp = new TopicPartition("topic1", 0);
        Map<TopicPartition, OffsetAndTimestamp> expectedResult = Collections.singletonMap(tp, null);
        Map<TopicPartition, Long> timestampToSearch = Collections.singletonMap(tp, 5L);
        Map<TopicPartition, OffsetAndTimestamp> result =
            assertDoesNotThrow(() -> consumer.offsetsForTimes(timestampToSearch, Duration.ZERO));
        assertEquals(expectedResult, result);
        verify(applicationEventHandler, never()).addAndGet(ArgumentMatchers.isA(ListOffsetsEvent.class));
    }

    @Test
    public void testWakeupCommitted() {
        consumer = newConsumer();
        final Map<TopicPartition, OffsetAndMetadata> offsets = mockTopicPartitionOffset();
        doAnswer(invocation -> {
            CompletableApplicationEvent<?> event = invocation.getArgument(0);
            assertInstanceOf(FetchCommittedOffsetsEvent.class, event);
            assertTrue(event.future().isCompletedExceptionally());
            return ConsumerUtils.getResult(event.future());
        })
            .when(applicationEventHandler)
            .addAndGet(any(FetchCommittedOffsetsEvent.class));

        consumer.wakeup();
        assertThrows(WakeupException.class, () -> consumer.committed(offsets.keySet()));
        assertNull(consumer.wakeupTrigger().getPendingTask());
    }

    @Test
    public void testNoWakeupInCloseCommit() {
        TopicPartition tp = new TopicPartition("topic1", 0);
        Properties props = requiredConsumerConfigAndGroupId("consumer-group");
        props.put(ConsumerConfig.ENABLE_AUTO_COMMIT_CONFIG, true);
        consumer = newConsumer(props);
        completeAssignmentChangeEventSuccessfully();
        consumer.assign(Collections.singleton(tp));
        doReturn(LeaderAndEpoch.noLeaderOrEpoch()).when(metadata).currentLeader(any());
        completeSeekUnvalidatedEventSuccessfully();
        consumer.seek(tp, 10);
        consumer.wakeup();

        AtomicReference<SyncCommitEvent> capturedEvent = new AtomicReference<>();
        doAnswer(invocation -> {
            ApplicationEvent event = invocation.getArgument(0);
            if (event instanceof SyncCommitEvent) {
                capturedEvent.set((SyncCommitEvent) event);
                ((SyncCommitEvent) event).markOffsetsReady();
            }
            return null;
        }).when(applicationEventHandler).add(any());
        completeUnsubscribeApplicationEventSuccessfully();
        consumer.close(CloseOptions.timeout(Duration.ZERO));

        // A commit was triggered and not completed exceptionally by the wakeup
        assertNotNull(capturedEvent.get());
        assertFalse(capturedEvent.get().future().isCompletedExceptionally());
    }

    @Test
    public void testCloseAwaitPendingAsyncCommitIncomplete() {
        time = new MockTime(1);
        consumer = newConsumer();

        // Commit async (incomplete)
        doReturn(LeaderAndEpoch.noLeaderOrEpoch()).when(metadata).currentLeader(any());
        final TopicPartition tp = new TopicPartition("foo", 0);
        completeAssignmentChangeEventSuccessfully();
        consumer.assign(Collections.singleton(tp));
        completeSeekUnvalidatedEventSuccessfully();
        consumer.seek(tp, 20);

        markOffsetsReadyForCommitEvent();
        consumer.commitAsync();

        Exception e = assertThrows(KafkaException.class, () -> consumer.close(CloseOptions.timeout(Duration.ofMillis(10))));
        assertInstanceOf(TimeoutException.class, e.getCause());
    }

    @Test
    public void testCloseAwaitPendingAsyncCommitComplete() {
        time = new MockTime(1);
        consumer = newConsumer();
        MockCommitCallback cb = new MockCommitCallback();

        // Commit async (complete)
        doReturn(LeaderAndEpoch.noLeaderOrEpoch()).when(metadata).currentLeader(any());
        final TopicPartition tp = new TopicPartition("foo", 0);
        completeAssignmentChangeEventSuccessfully();
        consumer.assign(Collections.singleton(tp));
        completeSeekUnvalidatedEventSuccessfully();
        consumer.seek(tp, 20);
        completeCommitAsyncApplicationEventSuccessfully();
        consumer.commitAsync(cb);

        completeUnsubscribeApplicationEventSuccessfully();
        assertDoesNotThrow(() -> consumer.close(CloseOptions.timeout(Duration.ofMillis(10))));
        assertEquals(1, cb.invoked);
    }


    @Test
    public void testInterceptorAutoCommitOnClose() {
        Properties props = requiredConsumerConfigAndGroupId("test-id");
        props.setProperty(ConsumerConfig.INTERCEPTOR_CLASSES_CONFIG, MockConsumerInterceptor.class.getName());
        props.setProperty(ConsumerConfig.ENABLE_AUTO_COMMIT_CONFIG, "true");

        consumer = newConsumer(props);
        assertEquals(1, MockConsumerInterceptor.INIT_COUNT.get());
        completeCommitSyncApplicationEventSuccessfully();
        completeUnsubscribeApplicationEventSuccessfully();

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

        assertEquals(1, MockConsumerInterceptor.ON_COMMIT_COUNT.get());
        assertEquals(1, MockConsumerInterceptor.CLOSE_COUNT.get());
    }

    @Test
    public void testInterceptorCommitSync() {
        Properties props = requiredConsumerConfigAndGroupId("test-id");
        props.setProperty(ConsumerConfig.INTERCEPTOR_CLASSES_CONFIG, MockConsumerInterceptor.class.getName());
        props.setProperty(ConsumerConfig.ENABLE_AUTO_COMMIT_CONFIG, "false");

        consumer = newConsumer(props);
        assertEquals(1, MockConsumerInterceptor.INIT_COUNT.get());
        completeCommitSyncApplicationEventSuccessfully();

        consumer.commitSync(mockTopicPartitionOffset());

        assertEquals(1, MockConsumerInterceptor.ON_COMMIT_COUNT.get());
    }

    @Test
    public void testNoInterceptorCommitSyncFailed() {
        Properties props = requiredConsumerConfigAndGroupId("test-id");
        props.setProperty(ConsumerConfig.INTERCEPTOR_CLASSES_CONFIG, MockConsumerInterceptor.class.getName());
        props.setProperty(ConsumerConfig.ENABLE_AUTO_COMMIT_CONFIG, "false");

        consumer = newConsumer(props);
        assertEquals(1, MockConsumerInterceptor.INIT_COUNT.get());
        KafkaException expected = new KafkaException("Test exception");
        completeCommitSyncApplicationEventExceptionally(expected);

        KafkaException actual = assertThrows(KafkaException.class, () -> consumer.commitSync(mockTopicPartitionOffset()));
        assertEquals(expected, actual);
        assertEquals(0, MockConsumerInterceptor.ON_COMMIT_COUNT.get());
    }

    @Test
    public void testInterceptorCommitAsync() {
        Properties props = requiredConsumerConfigAndGroupId("test-id");
        props.setProperty(ConsumerConfig.INTERCEPTOR_CLASSES_CONFIG, MockConsumerInterceptor.class.getName());
        props.setProperty(ConsumerConfig.ENABLE_AUTO_COMMIT_CONFIG, "false");

        consumer = newConsumer(props);
        assertEquals(1, MockConsumerInterceptor.INIT_COUNT.get());

        completeCommitAsyncApplicationEventSuccessfully();
        consumer.commitAsync(mockTopicPartitionOffset(), new MockCommitCallback());
        assertEquals(0, MockConsumerInterceptor.ON_COMMIT_COUNT.get());

        forceCommitCallbackInvocation();
        assertEquals(1, MockConsumerInterceptor.ON_COMMIT_COUNT.get());
    }

    @Test
    public void testNoInterceptorCommitAsyncFailed() {
        Properties props = requiredConsumerConfigAndGroupId("test-id");
        props.setProperty(ConsumerConfig.INTERCEPTOR_CLASSES_CONFIG, MockConsumerInterceptor.class.getName());
        props.setProperty(ConsumerConfig.ENABLE_AUTO_COMMIT_CONFIG, "false");

        consumer = newConsumer(props);
        assertEquals(1, MockConsumerInterceptor.INIT_COUNT.get());
        completeCommitAsyncApplicationEventExceptionally(new KafkaException("Test exception"));

        consumer.commitAsync(mockTopicPartitionOffset(), new MockCommitCallback());
        assertEquals(0, MockConsumerInterceptor.ON_COMMIT_COUNT.get());

        forceCommitCallbackInvocation();
        assertEquals(0, MockConsumerInterceptor.ON_COMMIT_COUNT.get());
    }

    @Test
    public void testSubscribeGeneratesEvent() {
        consumer = newConsumer();
        String topic = "topic1";
        completeTopicSubscriptionChangeEventSuccessfully();
        consumer.subscribe(singletonList(topic));
        assertEquals(singleton(topic), consumer.subscription());
        assertTrue(consumer.assignment().isEmpty());
        verify(applicationEventHandler).addAndGet(ArgumentMatchers.isA(TopicSubscriptionChangeEvent.class));
    }

    @Test
    public void testSubscribePatternGeneratesEvent() {
        consumer = newConsumer();
        Pattern pattern = Pattern.compile("topic.*");
        completeTopicPatternSubscriptionChangeEventSuccessfully();
        consumer.subscribe(pattern);
        verify(applicationEventHandler).addAndGet(ArgumentMatchers.isA(TopicPatternSubscriptionChangeEvent.class));
    }

    @Test
    public void testUnsubscribeGeneratesUnsubscribeEvent() {
        consumer = newConsumer();
        completeUnsubscribeApplicationEventSuccessfully();

        consumer.unsubscribe();

        assertTrue(consumer.subscription().isEmpty());
        assertTrue(consumer.assignment().isEmpty());
        ArgumentCaptor<UnsubscribeEvent> eventCaptor = ArgumentCaptor.forClass(UnsubscribeEvent.class);
        verify(applicationEventHandler).add(eventCaptor.capture());

        // check the deadline is set to the default API timeout
        long deadline = time.milliseconds() + (int) ConsumerConfig.configDef().defaultValues().get(ConsumerConfig.DEFAULT_API_TIMEOUT_MS_CONFIG);
        assertTrue(eventCaptor.getValue().deadlineMs() <= deadline);
    }

    @Test
    public void testSubscribeToEmptyListActsAsUnsubscribe() {
        consumer = newConsumer();
        completeUnsubscribeApplicationEventSuccessfully();

        consumer.subscribe(Collections.emptyList());
        assertTrue(consumer.subscription().isEmpty());
        assertTrue(consumer.assignment().isEmpty());
        verify(applicationEventHandler).add(ArgumentMatchers.isA(UnsubscribeEvent.class));
    }

    @Test
    public void testSubscribeToNullTopicCollection() {
        consumer = newConsumer();
        assertThrows(IllegalArgumentException.class, () -> consumer.subscribe((List<String>) null));
    }

    @Test
    public void testSubscriptionOnNullTopic() {
        consumer = newConsumer();
        assertThrows(IllegalArgumentException.class, () -> consumer.subscribe(singletonList(null)));
    }

    @Test
    public void testSubscriptionOnEmptyTopic() {
        consumer = newConsumer();
        String emptyTopic = "  ";
        assertThrows(IllegalArgumentException.class, () -> consumer.subscribe(singletonList(emptyTopic)));
    }

    @Test
    public void testGroupMetadataAfterCreationWithGroupIdIsNull() {
        final Properties props = requiredConsumerConfig();
        final ConsumerConfig config = new ConsumerConfig(props);
        consumer = newConsumer(config);

        assertFalse(config.unused().contains(ConsumerConfig.AUTO_COMMIT_INTERVAL_MS_CONFIG));
        assertFalse(config.unused().contains(THROW_ON_FETCH_STABLE_OFFSET_UNSUPPORTED));
        final Throwable exception = assertThrows(InvalidGroupIdException.class, consumer::groupMetadata);
        assertEquals(
            "To use the group management or offset commit APIs, you must " +
                "provide a valid " + ConsumerConfig.GROUP_ID_CONFIG + " in the consumer configuration.",
            exception.getMessage()
        );
    }

    @Test
    public void testGroupMetadataAfterCreationWithGroupIdIsNotNull() {
        final String groupId = "consumerGroupA";
        consumer = newConsumer(requiredConsumerConfigAndGroupId(groupId));

        final ConsumerGroupMetadata groupMetadata = consumer.groupMetadata();

        assertEquals(groupId, groupMetadata.groupId());
        assertEquals(Optional.empty(), groupMetadata.groupInstanceId());
        assertEquals(JoinGroupRequest.UNKNOWN_GENERATION_ID, groupMetadata.generationId());
        assertEquals(JoinGroupRequest.UNKNOWN_MEMBER_ID, groupMetadata.memberId());
    }

    @Test
    public void testGroupMetadataAfterCreationWithGroupIdIsNotNullAndGroupInstanceIdSet() {
        final String groupId = "consumerGroupA";
        final String groupInstanceId = "groupInstanceId1";
        final Properties props = requiredConsumerConfigAndGroupId(groupId);
        props.put(ConsumerConfig.GROUP_INSTANCE_ID_CONFIG, groupInstanceId);
        consumer = newConsumer(props);

        final ConsumerGroupMetadata groupMetadata = consumer.groupMetadata();

        assertEquals(groupId, groupMetadata.groupId());
        assertEquals(Optional.of(groupInstanceId), groupMetadata.groupInstanceId());
        assertEquals(JoinGroupRequest.UNKNOWN_GENERATION_ID, groupMetadata.generationId());
        assertEquals(JoinGroupRequest.UNKNOWN_MEMBER_ID, groupMetadata.memberId());
    }

    private MemberStateListener captureGroupMetadataUpdateListener(final MockedStatic<RequestManagers> requestManagers) {
        ArgumentCaptor<MemberStateListener> applicationThreadMemberStateListener = ArgumentCaptor.forClass(MemberStateListener.class);
        requestManagers.verify(() -> RequestManagers.supplier(
            any(),
            any(),
            any(),
            any(),
            any(),
            any(),
            any(),
            any(),
            any(),
            any(),
            any(),
            any(),
            any(),
            any(),
            applicationThreadMemberStateListener.capture(),
            any(),
            any()
        ));
        return applicationThreadMemberStateListener.getValue();
    }

    @Test
    public void testGroupMetadataUpdate() {
        final String groupId = "consumerGroupA";
        try (final MockedStatic<RequestManagers> requestManagers = mockStatic(RequestManagers.class)) {
            consumer = newConsumer(requiredConsumerConfigAndGroupId(groupId));
            final ConsumerGroupMetadata oldGroupMetadata = consumer.groupMetadata();
            final MemberStateListener groupMetadataUpdateListener = captureGroupMetadataUpdateListener(requestManagers);
            final int expectedMemberEpoch = 42;
            final String expectedMemberId = "memberId";
            groupMetadataUpdateListener.onMemberEpochUpdated(
                Optional.of(expectedMemberEpoch),
                expectedMemberId
            );
            final ConsumerGroupMetadata newGroupMetadata = consumer.groupMetadata();
            assertEquals(oldGroupMetadata.groupId(), newGroupMetadata.groupId());
            assertEquals(expectedMemberId, newGroupMetadata.memberId());
            assertEquals(expectedMemberEpoch, newGroupMetadata.generationId());
            assertEquals(oldGroupMetadata.groupInstanceId(), newGroupMetadata.groupInstanceId());
        }
    }

    @SuppressWarnings("removal")
    @Test
    public void testGroupMetadataIsResetAfterUnsubscribe() {
        final String groupId = "consumerGroupA";
        try (final MockedStatic<RequestManagers> requestManagers = mockStatic(RequestManagers.class)) {
            consumer = newConsumer(requiredConsumerConfigAndGroupId(groupId));
            final MemberStateListener groupMetadataUpdateListener = captureGroupMetadataUpdateListener(requestManagers);
            consumer.subscribe(singletonList("topic"));
            final int memberEpoch = 42;
            final String memberId = "memberId";
            groupMetadataUpdateListener.onMemberEpochUpdated(Optional.of(memberEpoch), memberId);
            final ConsumerGroupMetadata groupMetadata = consumer.groupMetadata();
            assertNotEquals(JoinGroupRequest.UNKNOWN_GENERATION_ID, groupMetadata.generationId());
            assertNotEquals(JoinGroupRequest.UNKNOWN_MEMBER_ID, groupMetadata.memberId());
        }
        completeUnsubscribeApplicationEventSuccessfully();

        consumer.unsubscribe();

        final ConsumerGroupMetadata groupMetadataAfterUnsubscribe = new ConsumerGroupMetadata(
            groupId,
            JoinGroupRequest.UNKNOWN_GENERATION_ID,
            JoinGroupRequest.UNKNOWN_MEMBER_ID,
            Optional.empty()
        );
        assertEquals(groupMetadataAfterUnsubscribe, consumer.groupMetadata());
    }

    private Optional<StreamsRebalanceData> captureStreamRebalanceData(final MockedStatic<RequestManagers> requestManagers) {
        ArgumentCaptor<Optional<StreamsRebalanceData>> streamRebalanceData = ArgumentCaptor.forClass(Optional.class);
        requestManagers.verify(() -> RequestManagers.supplier(
            any(),
            any(),
            any(),
            any(),
            any(),
            any(),
            any(),
            any(),
            any(),
            any(),
            any(),
            any(),
            any(),
            any(),
            any(),
            streamRebalanceData.capture(),
            any()
        ));
        return streamRebalanceData.getValue();
    }

    @Test
    public void testEmptyStreamRebalanceData() {
        final String groupId = "consumerGroupA";
        try (final MockedStatic<RequestManagers> requestManagers = mockStatic(RequestManagers.class)) {
            consumer = newConsumer(requiredConsumerConfigAndGroupId(groupId));
            final Optional<StreamsRebalanceData> groupMetadataUpdateListener = captureStreamRebalanceData(requestManagers);
            assertTrue(groupMetadataUpdateListener.isEmpty());
        }
    }

    @Test
    public void testStreamRebalanceData() {
        final String groupId = "consumerGroupA";
        try (final MockedStatic<RequestManagers> requestManagers = mockStatic(RequestManagers.class)) {
            StreamsRebalanceData streamsRebalanceData = new StreamsRebalanceData(UUID.randomUUID(), Optional.empty(), Optional.empty(), Map.of(), Map.of());
            consumer = newConsumerWithStreamRebalanceData(requiredConsumerConfigAndGroupId(groupId), streamsRebalanceData);
            final Optional<StreamsRebalanceData> groupMetadataUpdateListener = captureStreamRebalanceData(requestManagers);
            assertTrue(groupMetadataUpdateListener.isPresent());
            assertEquals(streamsRebalanceData, groupMetadataUpdateListener.get());
        }
    }

    /**
     * Tests that the consumer correctly invokes the callbacks for {@link ConsumerRebalanceListener} that was
     * specified. We don't go through the full effort to emulate heartbeats and correct group management here. We're
     * simply exercising the background {@link EventProcessor} does the correct thing when
     * {@link AsyncKafkaConsumer#poll(Duration)} is called.
     *
     * Note that we test {@link ConsumerRebalanceListener} that throws errors in its different callbacks. Failed
     * callback execution does <em>not</em> immediately errors. Instead, those errors are forwarded to the
     * application event thread for the {@link ConsumerMembershipManager} to handle.
     */
    @ParameterizedTest
    @MethodSource("listenerCallbacksInvokeSource")
    public void testListenerCallbacksInvoke(List<ConsumerRebalanceListenerMethodName> methodNames,
                                            Optional<RuntimeException> revokedError,
                                            Optional<RuntimeException> assignedError,
                                            Optional<RuntimeException> lostError,
                                            int expectedRevokedCount,
                                            int expectedAssignedCount,
                                            int expectedLostCount,
                                            Optional<RuntimeException> expectedException
                                            ) {
        consumer = newConsumer();
        CounterConsumerRebalanceListener consumerRebalanceListener = new CounterConsumerRebalanceListener(
                revokedError,
                assignedError,
                lostError
        );
        doReturn(Fetch.empty()).when(fetchCollector).collectFetch(any(FetchBuffer.class));
        completeTopicSubscriptionChangeEventSuccessfully();
        consumer.subscribe(Collections.singletonList("topic"), consumerRebalanceListener);
        SortedSet<TopicPartition> partitions = Collections.emptySortedSet();

        for (ConsumerRebalanceListenerMethodName methodName : methodNames) {
            CompletableBackgroundEvent<Void> e;
            if (methodName == ON_PARTITIONS_ASSIGNED) {
                e = new PartitionsAssignedEvent(Set.of(), partitions);
            } else {
                e = new PartitionsRemovedEvent(methodName, partitions);
            }
            backgroundEventQueue.add(e);
        }

        completeAsyncPollEventSuccessfully();
        // This will trigger the background event queue to process our background event message.
        // If any error is happening inside the rebalance callbacks, we expect the first exception to be thrown from poll.
        if (expectedException.isPresent()) {
            Exception exception = assertThrows(expectedException.get().getClass(), () -> consumer.poll(Duration.ZERO));
            assertEquals(expectedException.get().getMessage(), exception.getMessage());
            assertEquals(expectedException.get().getCause(), exception.getCause());
        } else {
            assertDoesNotThrow(() -> consumer.poll(Duration.ZERO));
        }

        assertEquals(expectedRevokedCount, consumerRebalanceListener.revokedCount());
        assertEquals(expectedAssignedCount, consumerRebalanceListener.assignedCount());
        assertEquals(expectedLostCount, consumerRebalanceListener.lostCount());
    }

    private static Stream<Arguments> listenerCallbacksInvokeSource() {
        Optional<RuntimeException> empty = Optional.empty();
        Optional<RuntimeException> error = Optional.of(new RuntimeException("Intentional error"));
        Optional<RuntimeException> kafkaException = Optional.of(new KafkaException("Intentional error"));
        Optional<RuntimeException> wrappedException = Optional.of(new KafkaException("User rebalance callback throws an error", error.get()));

        return Stream.of(
            // Tests if we don't have an event, the listener doesn't get called.
            Arguments.of(Collections.emptyList(), empty, empty, empty, 0, 0, 0, empty),

            // Tests if we get an event for a revocation, that we invoke our listener.
            Arguments.of(Collections.singletonList(ON_PARTITIONS_REVOKED), empty, empty, empty, 1, 0, 0, empty),

            // Tests if we get an event for an assignment, that we invoke our listener.
            Arguments.of(Collections.singletonList(ON_PARTITIONS_ASSIGNED), empty, empty, empty, 0, 1, 0, empty),

            // Tests that we invoke our listener even if it encounters an exception.
            Arguments.of(Collections.singletonList(ON_PARTITIONS_LOST), empty, empty, empty, 0, 0, 1, empty),

            // Tests that we invoke our listener even if it encounters an exception.
            Arguments.of(Collections.singletonList(ON_PARTITIONS_REVOKED), error, empty, empty, 1, 0, 0, wrappedException),

            // Tests that we invoke our listener even if it encounters an exception.
            Arguments.of(Collections.singletonList(ON_PARTITIONS_ASSIGNED), empty, error, empty, 0, 1, 0, wrappedException),

            // Tests that we invoke our listener even if it encounters an exception.
            Arguments.of(Collections.singletonList(ON_PARTITIONS_LOST), empty, empty, error, 0, 0, 1, wrappedException),

            // Tests that we invoke our listener even if it encounters an exception. Special case to test that a kafka exception is not wrapped.
            Arguments.of(Collections.singletonList(ON_PARTITIONS_REVOKED), kafkaException, empty, empty, 1, 0, 0, kafkaException),
            Arguments.of(Collections.singletonList(ON_PARTITIONS_ASSIGNED), empty, kafkaException, empty, 0, 1, 0, kafkaException),
            Arguments.of(Collections.singletonList(ON_PARTITIONS_LOST), empty, empty, kafkaException, 0, 0, 1, kafkaException),

            // Tests if we get separate events for revocation and then assignment--AND our revocation throws an error--
            // we still invoke the listeners correctly and throw the error.
            Arguments.of(Arrays.asList(ON_PARTITIONS_REVOKED, ON_PARTITIONS_ASSIGNED), error, empty, empty, 1, 1, 0, wrappedException),

            // Tests if we get separate events for revocation and then assignment--AND both throws an error--
            // we still invoke the listeners correctly and throw the first error.
            Arguments.of(Arrays.asList(ON_PARTITIONS_REVOKED, ON_PARTITIONS_ASSIGNED), kafkaException, error, empty, 1, 1, 0, kafkaException)
        );
    }

    @Test
    public void testBackgroundError() {
        final String groupId = "consumerGroupA";
        consumer = newConsumer(requiredConsumerConfigAndGroupId(groupId));

        final KafkaException expectedException = new KafkaException("Nobody expects the Spanish Inquisition");
        final ErrorEvent errorEvent = new ErrorEvent(expectedException);
        backgroundEventQueue.add(errorEvent);
        completeAssignmentChangeEventSuccessfully();
        consumer.assign(singletonList(new TopicPartition("topic", 0)));
        completeAsyncPollEventSuccessfully();
        final KafkaException exception = assertThrows(KafkaException.class, () -> consumer.poll(Duration.ZERO));

        assertEquals(expectedException.getMessage(), exception.getMessage());
    }

    @Test
    public void testMultipleBackgroundErrors() {
        final String groupId = "consumerGroupA";
        consumer = newConsumer(requiredConsumerConfigAndGroupId(groupId));

        final KafkaException expectedException1 = new KafkaException("Nobody expects the Spanish Inquisition");
        final ErrorEvent errorEvent1 = new ErrorEvent(expectedException1);
        backgroundEventQueue.add(errorEvent1);
        final KafkaException expectedException2 = new KafkaException("Spam, Spam, Spam");
        final ErrorEvent errorEvent2 = new ErrorEvent(expectedException2);
        backgroundEventQueue.add(errorEvent2);
        completeAssignmentChangeEventSuccessfully();
        consumer.assign(singletonList(new TopicPartition("topic", 0)));
        completeAsyncPollEventSuccessfully();
        final KafkaException exception = assertThrows(KafkaException.class, () -> consumer.poll(Duration.ZERO));

        assertEquals(expectedException1.getMessage(), exception.getMessage());
        assertTrue(backgroundEventQueue.isEmpty());
    }

    @Test
    public void testGroupRemoteAssignorUnusedIfGroupIdUndefined() {
        final Properties props = requiredConsumerConfig();
        props.put(ConsumerConfig.GROUP_REMOTE_ASSIGNOR_CONFIG, "someAssignor");
        props.put(ConsumerConfig.GROUP_PROTOCOL_CONFIG, GroupProtocol.CONSUMER.name().toLowerCase(Locale.ROOT));
        final ConsumerConfig config = new ConsumerConfig(props);
        consumer = newConsumer(config);

        assertTrue(config.unused().contains(ConsumerConfig.GROUP_REMOTE_ASSIGNOR_CONFIG));
    }

    @Test
    public void testGroupRemoteAssignorInClassicProtocol() {
        final Properties props = requiredConsumerConfig();
        props.put(ConsumerConfig.GROUP_ID_CONFIG, "consumerGroupA");
        props.put(ConsumerConfig.GROUP_PROTOCOL_CONFIG, GroupProtocol.CLASSIC.name().toLowerCase(Locale.ROOT));
        props.put(ConsumerConfig.GROUP_REMOTE_ASSIGNOR_CONFIG, "someAssignor");
        assertThrows(ConfigException.class, () -> new ConsumerConfig(props));
    }

    @Test
    public void testGroupRemoteAssignorUsedInConsumerProtocol() {
        final Properties props = requiredConsumerConfig();
        props.put(ConsumerConfig.GROUP_ID_CONFIG, "consumerGroupA");
        props.put(ConsumerConfig.GROUP_PROTOCOL_CONFIG, GroupProtocol.CONSUMER.name().toLowerCase(Locale.ROOT));
        props.put(ConsumerConfig.GROUP_REMOTE_ASSIGNOR_CONFIG, "someAssignor");
        props.put(ConsumerConfig.ENABLE_AUTO_COMMIT_CONFIG, false);
        final ConsumerConfig config = new ConsumerConfig(props);
        consumer = newConsumer(config);

        assertFalse(config.unused().contains(ConsumerConfig.GROUP_REMOTE_ASSIGNOR_CONFIG));
    }

    @Test
    public void testGroupIdNull() {
        final Properties props = requiredConsumerConfig();
        props.put(ConsumerConfig.AUTO_COMMIT_INTERVAL_MS_CONFIG, 10000);
        props.put(THROW_ON_FETCH_STABLE_OFFSET_UNSUPPORTED, true);
        final ConsumerConfig config = new ConsumerConfig(props);
        consumer = newConsumer(config);

        assertFalse(config.unused().contains(ConsumerConfig.AUTO_COMMIT_INTERVAL_MS_CONFIG));
        assertFalse(config.unused().contains(THROW_ON_FETCH_STABLE_OFFSET_UNSUPPORTED));
    }

    @Test
    public void testGroupIdNotNullAndValid() {
        final Properties props = requiredConsumerConfigAndGroupId("consumerGroupA");
        props.put(ConsumerConfig.AUTO_COMMIT_INTERVAL_MS_CONFIG, 10000);
        props.put(THROW_ON_FETCH_STABLE_OFFSET_UNSUPPORTED, true);
        props.put(ConsumerConfig.ENABLE_AUTO_COMMIT_CONFIG, false);
        final ConsumerConfig config = new ConsumerConfig(props);
        consumer = newConsumer(config);

        assertTrue(config.unused().contains(ConsumerConfig.AUTO_COMMIT_INTERVAL_MS_CONFIG));
        assertTrue(config.unused().contains(THROW_ON_FETCH_STABLE_OFFSET_UNSUPPORTED));
    }

    @Test
    public void testEnsurePollEventSentOnConsumerPoll() {
        SubscriptionState subscriptions = new SubscriptionState(new LogContext(), AutoOffsetResetStrategy.NONE);
        consumer = newConsumer(
                mock(FetchBuffer.class),
                new ConsumerInterceptors<>(Collections.emptyList(), metrics),
                mock(ConsumerRebalanceListenerInvoker.class),
                subscriptions);
        final TopicPartition tp = new TopicPartition("topic", 0);
        final List<ConsumerRecord<String, String>> records = singletonList(
                new ConsumerRecord<>("topic", 0, 2, "key1", "value1"));
        doAnswer(invocation -> Fetch.forPartition(tp, records, true, new OffsetAndMetadata(3, Optional.of(0), "")))
                .when(fetchCollector)
                .collectFetch(Mockito.any(FetchBuffer.class));

        completeTopicSubscriptionChangeEventSuccessfully();
        consumer.subscribe(singletonList("topic1"));
        completeAsyncPollEventSuccessfully();
        consumer.poll(Duration.ofMillis(100));
        verify(applicationEventHandler, atLeastOnce()).add(any(AsyncPollEvent.class));
    }

    private Properties requiredConsumerConfigAndGroupId(final String groupId) {
        final Properties props = requiredConsumerConfig();
        props.put(ConsumerConfig.GROUP_ID_CONFIG, groupId);
        return props;
    }

    @Test
    public void testLongPollWaitIsLimited() {
        consumer = newConsumer();
        String topicName = "topic1";
        completeTopicSubscriptionChangeEventSuccessfully();
        consumer.subscribe(singletonList(topicName));

        assertEquals(singleton(topicName), consumer.subscription());
        assertTrue(consumer.assignment().isEmpty());

        final int partition = 3;
        final TopicPartition tp = new TopicPartition(topicName, partition);
        final List<ConsumerRecord<String, String>> records = asList(
            new ConsumerRecord<>(topicName, partition, 2, "key1", "value1"),
            new ConsumerRecord<>(topicName, partition, 3, "key2", "value2")
        );
        final OffsetAndMetadata nextOffsetAndMetadata = new OffsetAndMetadata(4, Optional.of(0), "");

        // On the first iteration, return no data; on the second, return two records
        Set<TopicPartition> partitions = singleton(tp);
        doAnswer(invocation -> {
            // Mock the subscription being assigned as the first fetch is collected
            consumer.subscriptions().assignFromSubscribed(partitions);
            consumer.setGroupAssignmentSnapshot(partitions);
            return Fetch.empty();
        }).doAnswer(invocation ->
            Fetch.forPartition(tp, records, true, nextOffsetAndMetadata)
        ).when(fetchCollector).collectFetch(any(FetchBuffer.class));

        completeAsyncPollEventSuccessfully();
        // And then poll for up to 10000ms, which should return 2 records without timing out
        ConsumerRecords<?, ?> returnedRecords = consumer.poll(Duration.ofMillis(10000));
        assertEquals(2, returnedRecords.count());
        assertEquals(4, returnedRecords.nextOffsets().get(tp).offset());
        assertEquals(Optional.of(0), returnedRecords.nextOffsets().get(tp).leaderEpoch());

        assertEquals(singleton(topicName), consumer.subscription());
        assertEquals(partitions, consumer.assignment());
    }

    /**
     * Verifies that with manual partition assignment, poll() blocks for
     * the full user-supplied timeout instead of spinning in a busy loop.
     * Before the fix, AbstractHeartbeatRequestManager.maximumTimeToWait() returned 0
     * while the membership state was UNSUBSCRIBED (the state used for manual assignment),
     * causing pollForFetches() to call fetchBuffer.awaitWakeup() with a 0-timeout timer
     * and re-enter the poll loop immediately. After KAFKA-20426, maximumTimeToWait() returns
     * Long.MAX_VALUE in that state so the application thread can block for the full timeout.
     */
    @Test
    public void testPollWithManualAssignmentDoesNotBusyLoop() {
        FetchBuffer fetchBuffer = mock(FetchBuffer.class);
        ConsumerInterceptors<String, String> interceptors = mock(ConsumerInterceptors.class);
        ConsumerRebalanceListenerInvoker rebalanceListenerInvoker = mock(ConsumerRebalanceListenerInvoker.class);
        SubscriptionState subscriptions = new SubscriptionState(new LogContext(), AutoOffsetResetStrategy.NONE);
        consumer = newConsumer(fetchBuffer, interceptors, rebalanceListenerInvoker, subscriptions);

        final TopicPartition tp = new TopicPartition("topic1", 0);

        // Manual assignment with valid position so pollForFetches() does not shrink pollTimeout to retryBackoffMs.
        subscriptions.assignFromUser(singleton(tp));
        subscriptions.seek(tp, 0);

        // Simulate the FIXED behavior of AbstractHeartbeatRequestManager#maximumTimeToWait() when the
        // membership state is UNSUBSCRIBED, i.e. user called assign().
        doReturn(Long.MAX_VALUE).when(applicationEventHandler).maximumTimeToWait();

        doReturn(Fetch.empty()).when(fetchCollector).collectFetch(any(FetchBuffer.class));
        doReturn(LeaderAndEpoch.noLeaderOrEpoch()).when(metadata).currentLeader(any());

        // Capture the Timer passed to awaitWakeup so we can assert it was given the full
        // poll timeout, i.e. no busy loop. Also advance mock time by the timer's remaining ms so the
        // outer poll timer expires after a single wait.
        AtomicReference<Long> awaitTimerInitialMs = new AtomicReference<>();
        doAnswer(invocation -> {
            Timer pollTimer = invocation.getArgument(0, Timer.class);
            awaitTimerInitialMs.compareAndSet(null, pollTimer.remainingMs());
            time.sleep(pollTimer.remainingMs());
            pollTimer.update();
            return null;
        }).when(fetchBuffer).awaitWakeup(any(Timer.class));

        final long pollTimeoutMs = 500;
        consumer.poll(Duration.ofMillis(pollTimeoutMs));

        // The poll timer passed to awaitWakeup must have been set to the full user timeout,
        // proving pollTimeout was NOT clamped to 0 (busy loop) by maximumTimeToWait().
        assertNotNull(awaitTimerInitialMs.get(), "fetchBuffer.awaitWakeup was never called");
        assertEquals(pollTimeoutMs, awaitTimerInitialMs.get(),
            "Expected poll wait timer to use the full user timeout (no busy loop), but was " + awaitTimerInitialMs.get());

        // Only a single wait cycle should have happened
        verify(fetchBuffer, times(1)).awaitWakeup(any(Timer.class));
    }

    /**
     * Tests {@link AsyncKafkaConsumer#processBackgroundEvents(Future, Timer, Predicate, boolean) processBackgroundEvents}
     * handles the case where the {@link Future} takes a bit of time to complete, but does within the timeout.
     */
    @Test
    public void testProcessBackgroundEventsWithInitialDelay() throws Exception {
        consumer = newConsumer();
        Timer timer = time.timer(1000);
        CompletableFuture<?> future = mock(CompletableFuture.class);
        CountDownLatch latch = new CountDownLatch(3);

        // Mock our call to Future.get(timeout) so that it mimics a delay of 200 milliseconds. Keep in mind that
        // the incremental timeout inside processBackgroundEvents is 100 seconds for each pass. Our first two passes
        // will exceed the incremental timeout, but the third will return.
        doAnswer(invocation -> {
            latch.countDown();

            if (latch.getCount() > 0) {
                long timeout = invocation.getArgument(0, Long.class);
                timer.sleep(timeout);
                throw new java.util.concurrent.TimeoutException("Intentional timeout");
            }

            future.complete(null);
            return null;
        }).when(future).get(any(Long.class), any(TimeUnit.class));

        consumer.processBackgroundEvents(future, timer, e -> false, false);

        // 800 is the 1000 ms timeout (above) minus the 200 ms delay for the two incremental timeouts/retries.
        assertEquals(800, timer.remainingMs());
    }

    /**
     * Tests {@link AsyncKafkaConsumer#processBackgroundEvents(Future, Timer, Predicate, boolean) processBackgroundEvents}
     * handles the case where the {@link Future} is already complete when invoked, so it doesn't have to wait.
     */
    @Test
    public void testProcessBackgroundEventsWithoutDelay() {
        consumer = newConsumer();
        Timer timer = time.timer(1000);

        // Create a future that is already completed.
        CompletableFuture<?> future = CompletableFuture.completedFuture(null);

        consumer.processBackgroundEvents(future, timer, e -> false, false);

        // Because we didn't need to perform a timed get, we should still have every last millisecond
        // of our initial timeout.
        assertEquals(1000, timer.remainingMs());
    }

    /**
     * Tests {@link AsyncKafkaConsumer#processBackgroundEvents(Future, Timer, Predicate, boolean) processBackgroundEvents}
     * handles the case where the {@link Future} does not complete within the timeout.
     */
    @Test
    public void testProcessBackgroundEventsTimesOut() throws Exception {
        consumer = newConsumer();
        Timer timer = time.timer(1000);
        CompletableFuture<?> future = mock(CompletableFuture.class);

        doAnswer(invocation -> {
            long timeout = invocation.getArgument(0, Long.class);
            timer.sleep(timeout);
            throw new java.util.concurrent.TimeoutException("Intentional timeout");
        }).when(future).get(any(Long.class), any(TimeUnit.class));

        assertThrows(TimeoutException.class, () -> consumer.processBackgroundEvents(future, timer, e -> false, false));

        // Because we forced our mocked future to continuously time out, we should have no time remaining.
        assertEquals(0, timer.remainingMs());
    }

    /**
     * Tests that calling {@link Thread#interrupt()} before {@link KafkaConsumer#poll(Duration)}
     * causes {@link InterruptException} to be thrown.
     */
    @Test
    public void testPollThrowsInterruptExceptionIfInterrupted() {
        consumer = newConsumer();
        final String topicName = "foo";
        final int partition = 3;
        final TopicPartition tp = new TopicPartition(topicName, partition);
        doReturn(Fetch.empty()).when(fetchCollector).collectFetch(any(FetchBuffer.class));
        doReturn(LeaderAndEpoch.noLeaderOrEpoch()).when(metadata).currentLeader(any());
        completeAssignmentChangeEventSuccessfully();
        consumer.assign(singleton(tp));

        // interrupt the thread and call poll
        try {
            Thread.currentThread().interrupt();
            completeAsyncPollEventSuccessfully();
            assertThrows(InterruptException.class, () -> consumer.poll(Duration.ZERO));
        } finally {
            // clear interrupted state again since this thread may be reused by JUnit
            Thread.interrupted();
        }
        assertDoesNotThrow(() -> consumer.poll(Duration.ZERO));
    }

    @Test
    void testReaperInvokedInClose() {
        consumer = newConsumer();
        completeUnsubscribeApplicationEventSuccessfully();
        consumer.close();
        verify(backgroundEventReaper).reap(backgroundEventQueue);
    }

    @Test
    void testReaperInvokedInUnsubscribe() {
        consumer = newConsumer();
        completeUnsubscribeApplicationEventSuccessfully();
        consumer.unsubscribe();
        verify(backgroundEventReaper).reap(time.milliseconds());
    }

    @Test
    void testReaperInvokedInPoll() {
        consumer = newConsumer();
        doReturn(Fetch.empty()).when(fetchCollector).collectFetch(any(FetchBuffer.class));
        completeTopicSubscriptionChangeEventSuccessfully();
        consumer.subscribe(Collections.singletonList("topic"));
        completeAsyncPollEventSuccessfully();
        consumer.poll(Duration.ZERO);
        verify(backgroundEventReaper).reap(time.milliseconds());
    }

    @Test
    public void testUnsubscribeWithoutGroupId() {
        consumer = newConsumerWithoutGroupId();

        completeUnsubscribeApplicationEventSuccessfully();
        consumer.unsubscribe();
        verify(applicationEventHandler).add(ArgumentMatchers.isA(UnsubscribeEvent.class));
    }

    private static Stream<CompletableBackgroundEvent<?>> assignmentEventsSource() {
        return Stream.of(
            new PartitionsAssignedEvent(Set.of(), new TreeSet<>(TOPIC_PARTITION_COMPARATOR)),
            new StreamsTasksAssignedEvent(
                new TreeSet<>(TOPIC_PARTITION_COMPARATOR),
                new TreeSet<>(TOPIC_PARTITION_COMPARATOR),
                new StreamsRebalanceData.Assignment(Set.of(), Set.of(), Set.of(), true))
        );
    }

    /**
     * Test to ensure that assignment updates are not applied while unsubscribing
     * (it would cause an IllegalArgumentException when calling unsubscribe()).
     * Validates the fix for KAFKA-20428.
     */
    @ParameterizedTest
    @MethodSource("assignmentEventsSource")
    public void testUnsubscribeWithPendingAssignmentEvent(CompletableBackgroundEvent<?> assignedEvent) {
        consumer = newConsumer(requiredConsumerConfigAndGroupId("consumerGroup"));
        completeTopicSubscriptionChangeEventSuccessfully();
        consumer.subscribe(singletonList("topic"));
        completeUnsubscribeApplicationEventSuccessfully();

        // Add assignment event to the background queue (simulating an ongoing reconciliation
        // that completed just before unsubscribe was called)
        backgroundEventQueue.add(assignedEvent);

        // The call to unsubscribe should complete successfully (assignment event not processed and completed exceptionally)
        assertDoesNotThrow(() -> consumer.unsubscribe());
        verify(applicationEventHandler, never().description("Reconciled assignment updates shouldn't be processed while unsubscribing"))
                .addAndGet(any(ApplyAssignmentEvent.class));
        assertTrue(assignedEvent.future().isCompletedExceptionally());
    }

    @Test
    public void testSeekToBeginning() {
        Collection<TopicPartition> topics = Collections.singleton(new TopicPartition("test", 0));
        consumer = newConsumer();
        consumer.seekToBeginning(topics);
        CompletableApplicationEvent<Void> event = addAndGetLastEnqueuedEvent();
        ResetOffsetEvent resetOffsetEvent = assertInstanceOf(ResetOffsetEvent.class, event);
        assertEquals(topics, new HashSet<>(resetOffsetEvent.topicPartitions()));
        assertEquals(AutoOffsetResetStrategy.EARLIEST, resetOffsetEvent.offsetResetStrategy());
    }

    @Test
    public void testSeekToBeginningWithException() {
        Collection<TopicPartition> topics = Collections.singleton(new TopicPartition("test", 0));
        consumer = newConsumer();
        completeResetOffsetEventExceptionally(new TimeoutException());
        assertThrows(TimeoutException.class, () -> consumer.seekToBeginning(topics));
    }

    @Test
    public void testSeekToEndWithException() {
        Collection<TopicPartition> topics = Collections.singleton(new TopicPartition("test", 0));
        consumer = newConsumer();
        completeResetOffsetEventExceptionally(new TimeoutException());
        assertThrows(TimeoutException.class, () -> consumer.seekToEnd(topics));
    }

    @Test
    public void testSeekToEnd() {
        Collection<TopicPartition> topics = Collections.singleton(new TopicPartition("test", 0));
        consumer = newConsumer();
        consumer.seekToEnd(topics);
        CompletableApplicationEvent<Void> event = addAndGetLastEnqueuedEvent();
        ResetOffsetEvent resetOffsetEvent = assertInstanceOf(ResetOffsetEvent.class, event);
        assertEquals(topics, new HashSet<>(resetOffsetEvent.topicPartitions()));
        assertEquals(AutoOffsetResetStrategy.LATEST, resetOffsetEvent.offsetResetStrategy());
    }

    @Test
    public void testSubscribeToRe2JPatternValidation() {
        consumer = newConsumer();

        Throwable t = assertThrows(IllegalArgumentException.class, () -> consumer.subscribe((SubscriptionPattern) null));
        assertEquals("Topic pattern to subscribe to cannot be null", t.getMessage());

        t = assertThrows(IllegalArgumentException.class, () -> consumer.subscribe(new SubscriptionPattern("")));
        assertEquals("Topic pattern to subscribe to cannot be empty", t.getMessage());

        assertDoesNotThrow(() -> consumer.subscribe(new SubscriptionPattern("t*")));

        assertThrows(IllegalArgumentException.class, () -> consumer.subscribe(new SubscriptionPattern("t*"), null));
        assertDoesNotThrow(() -> consumer.subscribe(new SubscriptionPattern("t*"), mock(ConsumerRebalanceListener.class)));
    }

    @Test
    public void testSubscribeToRe2JPatternThrowsIfNoGroupId() {
        consumer = newConsumer(requiredConsumerConfig());
        assertThrows(InvalidGroupIdException.class, () -> consumer.subscribe(new SubscriptionPattern("t*")));
        assertThrows(InvalidGroupIdException.class, () -> consumer.subscribe(new SubscriptionPattern("t*"),
            mock(ConsumerRebalanceListener.class)));
    }

    @Test
    public void testSubscribeToRe2JPatternGeneratesEvent() {
        consumer = newConsumer();
        completeTopicRe2JPatternSubscriptionChangeEventSuccessfully();

        consumer.subscribe(new SubscriptionPattern("t*"));
        verify(applicationEventHandler).addAndGet(ArgumentMatchers.isA(TopicRe2JPatternSubscriptionChangeEvent.class));

        clearInvocations(applicationEventHandler);
        consumer.subscribe(new SubscriptionPattern("t*"), mock(ConsumerRebalanceListener.class));
        verify(applicationEventHandler).addAndGet(ArgumentMatchers.isA(TopicRe2JPatternSubscriptionChangeEvent.class));
    }

    // SubscriptionPattern is supported as of ConsumerGroupHeartbeatRequest v1. Clients using subscribe
    // (SubscribePattern) against older broker versions should get UnsupportedVersionException on poll after subscribe
    @Test
    public void testSubscribePatternAgainstBrokerNotSupportingRegex() throws InterruptedException {
        final Properties props = requiredConsumerConfig();
        props.put(ConsumerConfig.GROUP_ID_CONFIG, "group-id");
        props.put(ConsumerConfig.ENABLE_AUTO_COMMIT_CONFIG, "false");

        final ConsumerConfig config = new ConsumerConfig(props);

        ConsumerMetadata metadata = new ConsumerMetadata(0, 0, Long.MAX_VALUE, false, false,
            mock(SubscriptionState.class), new LogContext(), new ClusterResourceListeners());
        MockClient client = new MockClient(time, metadata);
        MetadataResponse initialMetadata = RequestTestUtils.metadataUpdateWithIds(1, Map.of("topic1", 2),
            Map.of("topic1", Uuid.randomUuid()));
        client.updateMetadata(initialMetadata);
        // ConsumerGroupHeartbeat v0 does not support broker-side regex resolution
        client.setNodeApiVersions(NodeApiVersions.create(ApiKeys.CONSUMER_GROUP_HEARTBEAT.id, (short) 0, (short) 0));

        // Mock response to find coordinator
        Node node = metadata.fetch().nodes().get(0);
        client.prepareResponseFrom(FindCoordinatorResponse.prepareResponse(Errors.NONE, "group-id", node), node);

        // Mock HB response (needed so that the MockClient builds the request)
        ConsumerGroupHeartbeatResponse result =
            new ConsumerGroupHeartbeatResponse(new ConsumerGroupHeartbeatResponseData()
                .setMemberId("")
                .setMemberEpoch(0));
        Node coordinator = new Node(Integer.MAX_VALUE - node.id(), node.host(), node.port());
        client.prepareResponseFrom(result, coordinator);

        SubscriptionState subscriptionState = mock(SubscriptionState.class);

        consumer = new AsyncKafkaConsumer<>(
            new LogContext(),
            time,
            config,
            new StringDeserializer(),
            new StringDeserializer(),
            client,
            subscriptionState,
            metadata
        );
        completeTopicRe2JPatternSubscriptionChangeEventSuccessfully();

        SubscriptionPattern pattern = new SubscriptionPattern("t*");
        consumer.subscribe(pattern);
        when(subscriptionState.subscriptionPattern()).thenReturn(pattern);
        TestUtils.waitForCondition(() -> {
            try {
                // The request is generated in the background thread so allow for that
                // async operation to happen to detect the failure.
                consumer.poll(Duration.ZERO);
                return false;
            } catch (UnsupportedVersionException e) {
                return true;
            }
        }, "Consumer did not throw the expected UnsupportedVersionException on poll");
    }

    @Test
    public void testRecordBackgroundEventQueueSizeAndBackgroundEventQueueTime() {
        consumer = newConsumer(
                mock(FetchBuffer.class),
                mock(ConsumerInterceptors.class),
                mock(ConsumerRebalanceListenerInvoker.class),
                mock(SubscriptionState.class));
        Metrics metrics = consumer.metricsRegistry();
        AsyncConsumerMetrics asyncConsumerMetrics = consumer.asyncConsumerMetrics();

        PartitionsRemovedEvent event = new PartitionsRemovedEvent(ON_PARTITIONS_REVOKED, Collections.emptySortedSet());
        event.setEnqueuedMs(time.milliseconds());
        backgroundEventQueue.add(event);
        asyncConsumerMetrics.recordBackgroundEventQueueSize(1);

        time.sleep(10);
        consumer.processBackgroundEvents();
        assertEquals(0, (double) metrics.metric(metrics.metricName("background-event-queue-size", CONSUMER_METRIC_GROUP)).metricValue());
        assertEquals(10, (double) metrics.metric(metrics.metricName("background-event-queue-time-avg", CONSUMER_METRIC_GROUP)).metricValue());
        assertEquals(10, (double) metrics.metric(metrics.metricName("background-event-queue-time-max", CONSUMER_METRIC_GROUP)).metricValue());
    }

    @Test
    public void testFailConstructor() {
        final Properties props = requiredConsumerConfig();
        props.put(ConsumerConfig.GROUP_ID_CONFIG, "group-id");
        props.put(ConsumerConfig.METRIC_REPORTER_CLASSES_CONFIG, "an.invalid.class");
        final ConsumerConfig config = new ConsumerConfig(props);

        try (LogCaptureAppender appender = LogCaptureAppender.createAndRegister()) {
            KafkaException ce = assertThrows(
                KafkaException.class,
                () -> newConsumer(config));
            assertTrue(ce.getMessage().contains("Failed to construct kafka consumer"), "Unexpected exception message: " + ce.getMessage());
            assertTrue(ce.getCause().getMessage().contains("Class an.invalid.class cannot be found"), "Unexpected cause: " + ce.getCause());

            boolean npeLogged = appender.getEvents().stream()
                .flatMap(event -> event.getThrowableInfo().stream())
                .anyMatch(str -> str.contains("NullPointerException"));

            assertFalse(npeLogged, "Unexpected NullPointerException during consumer construction");
        }
    }

    private Map<TopicPartition, OffsetAndMetadata> mockTopicPartitionOffset() {
        final TopicPartition t0 = new TopicPartition("t0", 2);
        final TopicPartition t1 = new TopicPartition("t0", 3);
        Map<TopicPartition, OffsetAndMetadata> topicPartitionOffsets = new HashMap<>();
        topicPartitionOffsets.put(t0, new OffsetAndMetadata(10L));
        topicPartitionOffsets.put(t1, new OffsetAndMetadata(20L));
        return topicPartitionOffsets;
    }

    private Map<TopicPartition, OffsetAndTimestampInternal> mockOffsetAndTimestamp() {
        final TopicPartition t0 = new TopicPartition("t0", 2);
        final TopicPartition t1 = new TopicPartition("t0", 3);
        Map<TopicPartition, OffsetAndTimestampInternal> offsetAndTimestamp = new HashMap<>();
        offsetAndTimestamp.put(t0, new OffsetAndTimestampInternal(5L, 1L, Optional.empty()));
        offsetAndTimestamp.put(t1, new OffsetAndTimestampInternal(6L, 3L, Optional.empty()));
        return offsetAndTimestamp;
    }

    private Map<TopicPartition, Long> mockTimestampToSearch() {
        final TopicPartition t0 = new TopicPartition("t0", 2);
        final TopicPartition t1 = new TopicPartition("t0", 3);
        Map<TopicPartition, Long> timestampToSearch = new HashMap<>();
        timestampToSearch.put(t0, 1L);
        timestampToSearch.put(t1, 2L);
        return timestampToSearch;
    }

    private void completeCommitAsyncApplicationEventExceptionally(Exception ex) {
        doAnswer(invocation -> {
            AsyncCommitEvent event = invocation.getArgument(0);
            event.markOffsetsReady();
            event.future().completeExceptionally(ex);
            return null;
        }).when(applicationEventHandler).add(ArgumentMatchers.isA(AsyncCommitEvent.class));
    }

    private void completeCommitSyncApplicationEventExceptionally(Exception ex) {
        doAnswer(invocation -> {
            SyncCommitEvent event = invocation.getArgument(0);
            event.markOffsetsReady();
            event.future().completeExceptionally(ex);
            return null;
        }).when(applicationEventHandler).add(ArgumentMatchers.isA(SyncCommitEvent.class));
    }

    private void completeResetOffsetEventExceptionally(Exception ex) {
        doThrow(ex).when(applicationEventHandler).addAndGet(ArgumentMatchers.isA(ResetOffsetEvent.class));
    }

    private void completeCommitAsyncApplicationEventSuccessfully() {
        doAnswer(invocation -> {
            AsyncCommitEvent event = invocation.getArgument(0);
            event.markOffsetsReady();
            event.future().complete(null);
            return null;
        }).when(applicationEventHandler).add(ArgumentMatchers.isA(AsyncCommitEvent.class));
    }

    private void completeCommitSyncApplicationEventSuccessfully() {
        doAnswer(invocation -> {
            SyncCommitEvent event = invocation.getArgument(0);
            event.markOffsetsReady();
            event.future().complete(null);
            return null;
        }).when(applicationEventHandler).add(ArgumentMatchers.isA(SyncCommitEvent.class));
    }

    private void completeFetchedCommittedOffsetApplicationEventSuccessfully(final Map<TopicPartition, OffsetAndMetadata> committedOffsets) {
        doReturn(committedOffsets)
            .when(applicationEventHandler)
            .addAndGet(any(FetchCommittedOffsetsEvent.class));

        doAnswer(invocation -> {
            FetchCommittedOffsetsEvent event = invocation.getArgument(0);
            event.future().complete(committedOffsets);
            return null;
        }).when(applicationEventHandler).add(ArgumentMatchers.isA(FetchCommittedOffsetsEvent.class));
    }

    private void completeFetchedCommittedOffsetApplicationEventExceptionally(Exception ex) {
        doThrow(ex)
            .when(applicationEventHandler)
            .addAndGet(any(FetchCommittedOffsetsEvent.class));
    }

    private void completeUnsubscribeApplicationEventSuccessfully() {
        doAnswer(invocation -> {
            UnsubscribeEvent event = invocation.getArgument(0);
            consumer.subscriptions().unsubscribe();
            event.future().complete(null);
            return null;
        }).when(applicationEventHandler).add(ArgumentMatchers.isA(UnsubscribeEvent.class));
    }

    private void completeAssignmentChangeEventSuccessfully() {
        doAnswer(invocation -> {
            AssignmentChangeEvent event = invocation.getArgument(0);
            HashSet<TopicPartition> partitions = new HashSet<>(event.partitions());
            consumer.subscriptions().assignFromUser(partitions);
            event.future().complete(null);
            return null;
        }).when(applicationEventHandler).addAndGet(ArgumentMatchers.isA(AssignmentChangeEvent.class));
    }

    private void completeTopicSubscriptionChangeEventSuccessfully() {
        doAnswer(invocation -> {
            TopicSubscriptionChangeEvent event = invocation.getArgument(0);
            consumer.subscriptions().subscribe(event.topics(), event.listener());
            event.future().complete(null);
            return null;
        }).when(applicationEventHandler).addAndGet(ArgumentMatchers.isA(TopicSubscriptionChangeEvent.class));
    }

    private void completeTopicPatternSubscriptionChangeEventSuccessfully() {
        doAnswer(invocation -> {
            TopicPatternSubscriptionChangeEvent event = invocation.getArgument(0);
            consumer.subscriptions().subscribe(event.pattern(), event.listener());
            event.future().complete(null);
            return null;
        }).when(applicationEventHandler).addAndGet(ArgumentMatchers.isA(TopicPatternSubscriptionChangeEvent.class));
    }

    private void completeTopicRe2JPatternSubscriptionChangeEventSuccessfully() {
        doAnswer(invocation -> {
            TopicRe2JPatternSubscriptionChangeEvent event = invocation.getArgument(0);
            consumer.subscriptions().subscribe(event.pattern(), event.listener());
            event.future().complete(null);
            return null;
        }).when(applicationEventHandler).addAndGet(ArgumentMatchers.isA(TopicRe2JPatternSubscriptionChangeEvent.class));
    }

    private void completeSeekUnvalidatedEventSuccessfully() {
        doAnswer(invocation -> {
            SeekUnvalidatedEvent event = invocation.getArgument(0);
            SubscriptionState.FetchPosition newPosition = new SubscriptionState.FetchPosition(
                    event.offset(),
                    event.offsetEpoch(),
                    metadata.currentLeader(event.partition())
            );
            consumer.subscriptions().seekUnvalidated(event.partition(), newPosition);
            event.future().complete(null);
            return null;
        }).when(applicationEventHandler).addAndGet(ArgumentMatchers.isA(SeekUnvalidatedEvent.class));
    }

    private void forceCommitCallbackInvocation() {
        // Invokes callback
        consumer.commitAsync();
    }

    private void markOffsetsReadyForCommitEvent() {
        doAnswer(invocation -> {
            CommitEvent event = invocation.getArgument(0);
            event.markOffsetsReady();
            return null;
        }).when(applicationEventHandler).add(ArgumentMatchers.isA(CommitEvent.class));
    }

    @Test
    public void testCloseInvokesStreamsRebalanceListenerOnTasksRevokedWhenMemberEpochPositive() {
        final String groupId = "streamsGroup";
        final StreamsRebalanceData streamsRebalanceData = new StreamsRebalanceData(UUID.randomUUID(), Optional.empty(), Optional.empty(), Map.of(), Map.of());
        
        try (final MockedStatic<RequestManagers> requestManagers = mockStatic(RequestManagers.class)) {
            consumer = newConsumerWithStreamRebalanceData(requiredConsumerConfigAndGroupId(groupId), streamsRebalanceData);
            StreamsRebalanceListener mockStreamsListener = mock(StreamsRebalanceListener.class);
            consumer.subscribe(singletonList("topic"), mockStreamsListener);
            final MemberStateListener groupMetadataUpdateListener = captureGroupMetadataUpdateListener(requestManagers);
            final int memberEpoch = 42;
            final String memberId = "memberId";
            groupMetadataUpdateListener.onMemberEpochUpdated(Optional.of(memberEpoch), memberId);
            
            consumer.close(CloseOptions.timeout(Duration.ZERO));
            
            verify(mockStreamsListener).onTasksRevoked(any());
        }
    }
    
    @Test
    public void testCloseInvokesStreamsRebalanceListenerOnAllTasksLostWhenMemberEpochZeroOrNegative() {
        final String groupId = "streamsGroup";
        final StreamsRebalanceData streamsRebalanceData = new StreamsRebalanceData(UUID.randomUUID(), Optional.empty(), Optional.empty(), Map.of(), Map.of());
        
        try (final MockedStatic<RequestManagers> requestManagers = mockStatic(RequestManagers.class)) {
            consumer = newConsumerWithStreamRebalanceData(requiredConsumerConfigAndGroupId(groupId), streamsRebalanceData);
            StreamsRebalanceListener mockStreamsListener = mock(StreamsRebalanceListener.class);
            consumer.subscribe(singletonList("topic"), mockStreamsListener);
            final MemberStateListener groupMetadataUpdateListener = captureGroupMetadataUpdateListener(requestManagers);
            final int memberEpoch = 0;
            final String memberId = "memberId";
            groupMetadataUpdateListener.onMemberEpochUpdated(Optional.of(memberEpoch), memberId);
            
            consumer.close(CloseOptions.timeout(Duration.ZERO));
            
            verify(mockStreamsListener).onAllTasksLost();
        }
    }
    
    @Test
    public void testCloseWrapsStreamsRebalanceListenerException() {
        final String groupId = "streamsGroup";
        final StreamsRebalanceData streamsRebalanceData = new StreamsRebalanceData(UUID.randomUUID(), Optional.empty(), Optional.empty(), Map.of(), Map.of());
        
        try (final MockedStatic<RequestManagers> requestManagers = mockStatic(RequestManagers.class)) {
            consumer = newConsumerWithStreamRebalanceData(requiredConsumerConfigAndGroupId(groupId), streamsRebalanceData);
            StreamsRebalanceListener mockStreamsListener = mock(StreamsRebalanceListener.class);
            RuntimeException testException = new RuntimeException("Test streams listener exception");
            doThrow(testException).when(mockStreamsListener).onTasksRevoked(any());
            consumer.subscribe(singletonList("topic"), mockStreamsListener);
            final MemberStateListener groupMetadataUpdateListener = captureGroupMetadataUpdateListener(requestManagers);
            final int memberEpoch = 1;
            final String memberId = "memberId";
            groupMetadataUpdateListener.onMemberEpochUpdated(Optional.of(memberEpoch), memberId);
            
            KafkaException thrownException = assertThrows(KafkaException.class, 
                () -> consumer.close(CloseOptions.timeout(Duration.ZERO)));

            assertInstanceOf(RuntimeException.class, thrownException.getCause());
            assertTrue(thrownException.getCause().getMessage().contains("Test streams listener exception"));
            verify(mockStreamsListener).onTasksRevoked(any());
        }
    }

    /**
     * Tests that when the app thread fails to update the assignment as part of a reconciliation,
     * it sends an event with error to the background thread before throwing on poll.
     * This ensures the background reconciliation is completed exceptionally.
     */
    @Test
    public void testPartitionsAssignedEventSendsErrorWhenApplyAssignmentFails() {
        final InterruptException applyAssignmentError = new InterruptException("Thread was interrupted");

        consumer = newConsumer(requiredConsumerConfigAndGroupId("consumerGroup"));
        completeTopicSubscriptionChangeEventSuccessfully();
        consumer.subscribe(singletonList("topic"), new CounterConsumerRebalanceListener(
            Optional.empty(), Optional.empty(), Optional.empty()));

        // Make ApplyAssignmentEvent fail
        when(applicationEventHandler.addAndGet(any(ApplyAssignmentEvent.class)))
            .thenThrow(applyAssignmentError);

        // Add PartitionsAssignedEvent to background queue
        backgroundEventQueue.add(new PartitionsAssignedEvent(Set.of(), new TreeSet<>()));

        completeAsyncPollEventSuccessfully();

        // Poll should throw because it couldn't update the assignment to run callbacks
        assertSame(applyAssignmentError, assertThrows(InterruptException.class, () -> consumer.poll(Duration.ZERO)));

        // Verify that ConsumerRebalanceListenerCallbackCompletedEvent with error was sent
        ArgumentCaptor<ConsumerRebalanceListenerCallbackCompletedEvent> eventCaptor =
            ArgumentCaptor.forClass(ConsumerRebalanceListenerCallbackCompletedEvent.class);
        verify(applicationEventHandler).add(eventCaptor.capture());
        assertTrue(eventCaptor.getValue().error().isPresent());
        assertSame(applyAssignmentError, eventCaptor.getValue().error().get());
        assertEquals(ON_PARTITIONS_ASSIGNED, eventCaptor.getValue().methodName());
    }

    /**
     * Tests that when the app thread fails to update the assignment as part of a reconciliation,
     * it sends an event with error to the background thread before throwing on poll.
     * This ensures the background reconciliation is completed exceptionally.
     */
    @Test
    public void testStreamsTasksAssignedEventSendsErrorWhenApplyAssignmentFails() {
        final StreamsRebalanceData streamsRebalanceData = new StreamsRebalanceData(
            UUID.randomUUID(), Optional.empty(), Optional.empty(), Map.of(), Map.of());
        final InterruptException applyAssignmentError = new InterruptException("Thread was interrupted");

        consumer = newConsumerWithStreamRebalanceData(
            requiredConsumerConfigAndGroupId("streamsGroup"), streamsRebalanceData);
        completeTopicSubscriptionChangeEventSuccessfully();
        consumer.subscribe(singletonList("topic"), mock(StreamsRebalanceListener.class));

        // Make ApplyAssignmentEvent fail
        when(applicationEventHandler.addAndGet(any(ApplyAssignmentEvent.class)))
            .thenThrow(applyAssignmentError);

        // Add StreamsTasksAssignedEvent to background queue
        backgroundEventQueue.add(new StreamsTasksAssignedEvent(
            new TreeSet<>(), new TreeSet<>(),
            new StreamsRebalanceData.Assignment(Set.of(), Set.of(), Set.of(), true)));

        completeAsyncPollEventSuccessfully();

        // Poll should throw because it failed to update assignment to run callbacks
        assertSame(applyAssignmentError, assertThrows(InterruptException.class, () -> consumer.poll(Duration.ZERO)));

        // Verify that StreamsOnTasksAssignedCallbackCompletedEvent with error was sent
        ArgumentCaptor<StreamsOnTasksAssignedCallbackCompletedEvent> eventCaptor =
            ArgumentCaptor.forClass(StreamsOnTasksAssignedCallbackCompletedEvent.class);
        verify(applicationEventHandler).add(eventCaptor.capture());
        assertTrue(eventCaptor.getValue().error().isPresent());
        assertSame(applyAssignmentError, eventCaptor.getValue().error().get());
    }

    private void completeAsyncPollEventSuccessfully() {
        doAnswer(invocation -> {
            AsyncPollEvent event = invocation.getArgument(0);
            event.completeSuccessfully();
            return null;
        }).when(applicationEventHandler).add(ArgumentMatchers.isA(AsyncPollEvent.class));
    }
}
