/*
 * 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.flink.connector.kafka.source.enumerator;

import org.apache.flink.api.connector.source.Boundedness;
import org.apache.flink.api.connector.source.ReaderInfo;
import org.apache.flink.api.connector.source.mocks.MockSplitEnumeratorContext;
import org.apache.flink.connector.kafka.source.KafkaSourceOptions;
import org.apache.flink.connector.kafka.source.enumerator.initializer.NoStoppingOffsetsInitializer;
import org.apache.flink.connector.kafka.source.enumerator.initializer.OffsetsInitializer;
import org.apache.flink.connector.kafka.source.enumerator.subscriber.KafkaSubscriber;
import org.apache.flink.connector.kafka.source.split.KafkaPartitionSplit;
import org.apache.flink.connector.kafka.testutils.KafkaSourceTestEnv;
import org.apache.flink.mock.Whitebox;

import org.apache.kafka.clients.admin.AdminClient;
import org.apache.kafka.clients.admin.NewTopic;
import org.apache.kafka.clients.consumer.ConsumerConfig;
import org.apache.kafka.common.TopicPartition;
import org.apache.kafka.common.serialization.StringDeserializer;
import org.junit.AfterClass;
import org.junit.BeforeClass;
import org.junit.Test;

import java.util.ArrayList;
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.Map;
import java.util.Properties;
import java.util.Set;
import java.util.StringJoiner;
import java.util.regex.Pattern;

import static org.assertj.core.api.Assertions.assertThat;

/** Unit tests for {@link KafkaSourceEnumerator}. */
public class KafkaEnumeratorTest {
    private static final int NUM_SUBTASKS = 3;
    private static final String DYNAMIC_TOPIC_NAME = "dynamic_topic";
    private static final int NUM_PARTITIONS_DYNAMIC_TOPIC = 4;

    private static final String TOPIC1 = "topic";
    private static final String TOPIC2 = "pattern-topic";

    private static final int READER0 = 0;
    private static final int READER1 = 1;
    private static final Set<String> PRE_EXISTING_TOPICS =
            new HashSet<>(Arrays.asList(TOPIC1, TOPIC2));
    private static final int PARTITION_DISCOVERY_CALLABLE_INDEX = 0;
    private static final boolean ENABLE_PERIODIC_PARTITION_DISCOVERY = true;
    private static final boolean DISABLE_PERIODIC_PARTITION_DISCOVERY = false;
    private static final boolean INCLUDE_DYNAMIC_TOPIC = true;
    private static final boolean EXCLUDE_DYNAMIC_TOPIC = false;

    @BeforeClass
    public static void setup() throws Throwable {
        KafkaSourceTestEnv.setup();
        KafkaSourceTestEnv.setupTopic(TOPIC1, true, true, KafkaSourceTestEnv::getRecordsForTopic);
        KafkaSourceTestEnv.setupTopic(TOPIC2, true, true, KafkaSourceTestEnv::getRecordsForTopic);
    }

    @AfterClass
    public static void tearDown() throws Exception {
        KafkaSourceTestEnv.tearDown();
    }

    @Test
    public void testStartWithDiscoverPartitionsOnce() throws Exception {
        try (MockSplitEnumeratorContext<KafkaPartitionSplit> context =
                        new MockSplitEnumeratorContext<>(NUM_SUBTASKS);
                KafkaSourceEnumerator enumerator =
                        createEnumerator(context, DISABLE_PERIODIC_PARTITION_DISCOVERY)) {

            // Start the enumerator and it should schedule a one time task to discover and assign
            // partitions.
            enumerator.start();
            assertThat(context.getPeriodicCallables()).isEmpty();
            assertThat(context.getOneTimeCallables())
                    .as("A one time partition discovery callable should have been scheduled")
                    .hasSize(1);

            // enumerator just start noMoreNewPartitionSplits will be false
            assertThat((Boolean) Whitebox.getInternalState(enumerator, "noMoreNewPartitionSplits"))
                    .isFalse();
        }
    }

    @Test
    public void testStartWithPeriodicPartitionDiscovery() throws Exception {
        try (MockSplitEnumeratorContext<KafkaPartitionSplit> context =
                        new MockSplitEnumeratorContext<>(NUM_SUBTASKS);
                KafkaSourceEnumerator enumerator =
                        createEnumerator(context, ENABLE_PERIODIC_PARTITION_DISCOVERY)) {

            // Start the enumerator and it should schedule a one time task to discover and assign
            // partitions.
            enumerator.start();
            assertThat(context.getOneTimeCallables()).isEmpty();
            assertThat(context.getPeriodicCallables())
                    .as("A periodic partition discovery callable should have been scheduled")
                    .hasSize(1);

            // enumerator just start noMoreNewPartitionSplits will be false
            assertThat((Boolean) Whitebox.getInternalState(enumerator, "noMoreNewPartitionSplits"))
                    .isFalse();
        }
    }

    @Test
    public void testDiscoverPartitionsTriggersAssignments() throws Throwable {
        try (MockSplitEnumeratorContext<KafkaPartitionSplit> context =
                        new MockSplitEnumeratorContext<>(NUM_SUBTASKS);
                KafkaSourceEnumerator enumerator =
                        createEnumerator(context, DISABLE_PERIODIC_PARTITION_DISCOVERY)) {

            // Start the enumerator and it should schedule a one time task to discover and assign
            // partitions.
            enumerator.start();

            // register reader 0.
            registerReader(context, enumerator, READER0);
            registerReader(context, enumerator, READER1);
            assertThat(context.getSplitsAssignmentSequence()).isEmpty();

            // Run the partition discover callable and check the partition assignment.
            runOneTimePartitionDiscovery(context);

            // Verify assignments for reader 0.
            verifyLastReadersAssignments(
                    context, Arrays.asList(READER0, READER1), PRE_EXISTING_TOPICS, 1);
        }
    }

    @Test
    public void testReaderRegistrationTriggersAssignments() throws Throwable {
        try (MockSplitEnumeratorContext<KafkaPartitionSplit> context =
                        new MockSplitEnumeratorContext<>(NUM_SUBTASKS);
                KafkaSourceEnumerator enumerator =
                        createEnumerator(context, DISABLE_PERIODIC_PARTITION_DISCOVERY)) {

            // Start the enumerator and it should schedule a one time task to discover and assign
            // partitions.
            enumerator.start();
            runOneTimePartitionDiscovery(context);
            assertThat(context.getSplitsAssignmentSequence()).isEmpty();

            registerReader(context, enumerator, READER0);
            verifyLastReadersAssignments(
                    context, Collections.singleton(READER0), PRE_EXISTING_TOPICS, 1);

            registerReader(context, enumerator, READER1);
            verifyLastReadersAssignments(
                    context, Collections.singleton(READER1), PRE_EXISTING_TOPICS, 2);
        }
    }

    @Test
    public void testRunWithDiscoverPartitionsOnceToCheckNoMoreSplit() throws Throwable {
        try (MockSplitEnumeratorContext<KafkaPartitionSplit> context =
                        new MockSplitEnumeratorContext<>(NUM_SUBTASKS);
                KafkaSourceEnumerator enumerator =
                        createEnumerator(context, DISABLE_PERIODIC_PARTITION_DISCOVERY)) {

            // Start the enumerator and it should schedule a one time task to discover and assign
            // partitions.
            enumerator.start();
            assertThat(context.getOneTimeCallables())
                    .as("A one time partition discovery callable should have been scheduled")
                    .hasSize(1);
            assertThat(context.getPeriodicCallables()).isEmpty();
            // Run the partition discover callable and check the partition assignment.
            runOneTimePartitionDiscovery(context);

            // enumerator noMoreNewPartitionSplits first will be false, when execute
            // handlePartitionSplitChanges will be set true
            assertThat((Boolean) Whitebox.getInternalState(enumerator, "noMoreNewPartitionSplits"))
                    .isTrue();
        }
    }

    @Test
    public void testRunWithPeriodicPartitionDiscoveryOnceToCheckNoMoreSplit() throws Throwable {
        try (MockSplitEnumeratorContext<KafkaPartitionSplit> context =
                        new MockSplitEnumeratorContext<>(NUM_SUBTASKS);
                KafkaSourceEnumerator enumerator =
                        createEnumerator(context, ENABLE_PERIODIC_PARTITION_DISCOVERY)) {

            // Start the enumerator and it should schedule a one time task to discover and assign
            // partitions.
            enumerator.start();
            assertThat(context.getOneTimeCallables()).isEmpty();
            assertThat(context.getPeriodicCallables())
                    .as("A periodic partition discovery callable should have been scheduled")
                    .hasSize(1);
            // Run the partition discover callable and check the partition assignment.
            runPeriodicPartitionDiscovery(context);

            // enumerator noMoreNewPartitionSplits first will be false, even when execute
            // handlePartitionSplitChanges it still be false
            assertThat((Boolean) Whitebox.getInternalState(enumerator, "noMoreNewPartitionSplits"))
                    .isFalse();
        }
    }

    @Test
    public void testRunWithDiscoverPartitionsOnceWithZeroMsToCheckNoMoreSplit() throws Throwable {
        try (MockSplitEnumeratorContext<KafkaPartitionSplit> context =
                        new MockSplitEnumeratorContext<>(NUM_SUBTASKS);
                // set partitionDiscoveryIntervalMs = 0
                KafkaSourceEnumerator enumerator = createEnumerator(context, 0L)) {

            // Start the enumerator, and it should schedule a one time task to discover and assign
            // partitions.
            enumerator.start();
            assertThat(context.getOneTimeCallables())
                    .as("A one time partition discovery callable should have been scheduled")
                    .hasSize(1);
            assertThat(context.getPeriodicCallables()).isEmpty();
            // Run the partition discover callable and check the partition assignment.
            runOneTimePartitionDiscovery(context);

            // enumerator noMoreNewPartitionSplits first will be false, when execute
            // handlePartitionSplitChanges will be set true
            assertThat((Boolean) Whitebox.getInternalState(enumerator, "noMoreNewPartitionSplits"))
                    .isTrue();
        }
    }

    @Test(timeout = 30000L)
    public void testDiscoverPartitionsPeriodically() throws Throwable {
        try (MockSplitEnumeratorContext<KafkaPartitionSplit> context =
                        new MockSplitEnumeratorContext<>(NUM_SUBTASKS);
                KafkaSourceEnumerator enumerator =
                        createEnumerator(
                                context,
                                ENABLE_PERIODIC_PARTITION_DISCOVERY,
                                INCLUDE_DYNAMIC_TOPIC);
                AdminClient adminClient = KafkaSourceTestEnv.getAdminClient()) {

            startEnumeratorAndRegisterReaders(context, enumerator);

            // invoke partition discovery callable again and there should be no new assignments.
            runPeriodicPartitionDiscovery(context);
            assertThat(context.getSplitsAssignmentSequence())
                    .as("No assignments should be made because there is no partition change")
                    .hasSize(2);

            // create the dynamic topic.
            adminClient
                    .createTopics(
                            Collections.singleton(
                                    new NewTopic(
                                            DYNAMIC_TOPIC_NAME,
                                            NUM_PARTITIONS_DYNAMIC_TOPIC,
                                            (short) 1)))
                    .all()
                    .get();

            // invoke partition discovery callable again.
            while (true) {
                runPeriodicPartitionDiscovery(context);
                if (context.getSplitsAssignmentSequence().size() < 3) {
                    Thread.sleep(10);
                } else {
                    break;
                }
            }
            verifyLastReadersAssignments(
                    context,
                    Arrays.asList(READER0, READER1),
                    Collections.singleton(DYNAMIC_TOPIC_NAME),
                    3);
        } finally {
            try (AdminClient adminClient = KafkaSourceTestEnv.getAdminClient()) {
                adminClient.deleteTopics(Collections.singleton(DYNAMIC_TOPIC_NAME)).all().get();
            } catch (Exception e) {
                // Let it go.
            }
        }
    }

    @Test
    public void testAddSplitsBack() throws Throwable {
        try (MockSplitEnumeratorContext<KafkaPartitionSplit> context =
                        new MockSplitEnumeratorContext<>(NUM_SUBTASKS);
                KafkaSourceEnumerator enumerator =
                        createEnumerator(context, ENABLE_PERIODIC_PARTITION_DISCOVERY)) {

            startEnumeratorAndRegisterReaders(context, enumerator);

            // Simulate a reader failure.
            context.unregisterReader(READER0);
            enumerator.addSplitsBack(
                    context.getSplitsAssignmentSequence().get(0).assignment().get(READER0),
                    READER0);
            assertThat(context.getSplitsAssignmentSequence())
                    .as("The added back splits should have not been assigned")
                    .hasSize(2);

            // Simulate a reader recovery.
            registerReader(context, enumerator, READER0);
            verifyLastReadersAssignments(
                    context, Collections.singleton(READER0), PRE_EXISTING_TOPICS, 3);
        }
    }

    @Test
    public void testWorkWithPreexistingAssignments() throws Throwable {
        Set<TopicPartition> preexistingAssignments;
        try (MockSplitEnumeratorContext<KafkaPartitionSplit> context1 =
                        new MockSplitEnumeratorContext<>(NUM_SUBTASKS);
                KafkaSourceEnumerator enumerator =
                        createEnumerator(context1, ENABLE_PERIODIC_PARTITION_DISCOVERY)) {
            startEnumeratorAndRegisterReaders(context1, enumerator);
            preexistingAssignments =
                    asEnumState(context1.getSplitsAssignmentSequence().get(0).assignment());
        }

        try (MockSplitEnumeratorContext<KafkaPartitionSplit> context2 =
                        new MockSplitEnumeratorContext<>(NUM_SUBTASKS);
                KafkaSourceEnumerator enumerator =
                        createEnumerator(
                                context2,
                                ENABLE_PERIODIC_PARTITION_DISCOVERY ? 1 : -1,
                                PRE_EXISTING_TOPICS,
                                preexistingAssignments,
                                new Properties())) {
            enumerator.start();
            runPeriodicPartitionDiscovery(context2);

            registerReader(context2, enumerator, READER0);
            assertThat(context2.getSplitsAssignmentSequence()).isEmpty();

            registerReader(context2, enumerator, READER1);
            verifyLastReadersAssignments(
                    context2, Collections.singleton(READER1), PRE_EXISTING_TOPICS, 1);
        }
    }

    @Test
    public void testKafkaClientProperties() throws Exception {
        Properties properties = new Properties();
        String clientIdPrefix = "test-prefix";
        Integer defaultTimeoutMs = 99999;
        properties.setProperty(KafkaSourceOptions.CLIENT_ID_PREFIX.key(), clientIdPrefix);
        properties.setProperty(
                ConsumerConfig.REQUEST_TIMEOUT_MS_CONFIG, String.valueOf(defaultTimeoutMs));
        try (MockSplitEnumeratorContext<KafkaPartitionSplit> context =
                        new MockSplitEnumeratorContext<>(NUM_SUBTASKS);
                KafkaSourceEnumerator enumerator =
                        createEnumerator(
                                context,
                                ENABLE_PERIODIC_PARTITION_DISCOVERY ? 1 : -1,
                                PRE_EXISTING_TOPICS,
                                Collections.emptySet(),
                                properties)) {
            enumerator.start();

            AdminClient adminClient =
                    (AdminClient) Whitebox.getInternalState(enumerator, "adminClient");
            assertThat(adminClient).isNotNull();
            String clientId = (String) Whitebox.getInternalState(adminClient, "clientId");
            assertThat(clientId).isNotNull().startsWith(clientIdPrefix);
            assertThat(Whitebox.getInternalState(adminClient, "defaultApiTimeoutMs"))
                    .isEqualTo(defaultTimeoutMs);

            assertThat(clientId).isNotNull().startsWith(clientIdPrefix);
        }
    }

    @Test
    public void testSnapshotState() throws Throwable {
        try (MockSplitEnumeratorContext<KafkaPartitionSplit> context =
                        new MockSplitEnumeratorContext<>(NUM_SUBTASKS);
                KafkaSourceEnumerator enumerator = createEnumerator(context, false)) {
            enumerator.start();

            // No reader is registered, so the state should be empty
            final KafkaSourceEnumState state1 = enumerator.snapshotState(1L);
            assertThat(state1.assignedPartitions()).isEmpty();

            registerReader(context, enumerator, READER0);
            registerReader(context, enumerator, READER1);
            runOneTimePartitionDiscovery(context);

            // The state should contain splits assigned to READER0 and READER1
            final KafkaSourceEnumState state2 = enumerator.snapshotState(1L);
            verifySplitAssignmentWithPartitions(
                    getExpectedAssignments(
                            new HashSet<>(Arrays.asList(READER0, READER1)), PRE_EXISTING_TOPICS),
                    state2.assignedPartitions());
        }
    }

    @Test
    public void testPartitionChangeChecking() throws Throwable {
        try (MockSplitEnumeratorContext<KafkaPartitionSplit> context =
                        new MockSplitEnumeratorContext<>(NUM_SUBTASKS);
                KafkaSourceEnumerator enumerator =
                        createEnumerator(context, DISABLE_PERIODIC_PARTITION_DISCOVERY)) {
            enumerator.start();
            runOneTimePartitionDiscovery(context);
            registerReader(context, enumerator, READER0);
            verifyLastReadersAssignments(
                    context, Collections.singleton(READER0), PRE_EXISTING_TOPICS, 1);

            // All partitions of TOPIC1 and TOPIC2 should have been discovered now

            // Check partition change using only DYNAMIC_TOPIC_NAME-0
            TopicPartition newPartition = new TopicPartition(DYNAMIC_TOPIC_NAME, 0);
            Set<TopicPartition> fetchedPartitions = new HashSet<>();
            fetchedPartitions.add(newPartition);
            final KafkaSourceEnumerator.PartitionChange partitionChange =
                    enumerator.getPartitionChange(fetchedPartitions);

            // Since enumerator never met DYNAMIC_TOPIC_NAME-0, it should be mark as a new partition
            Set<TopicPartition> expectedNewPartitions = Collections.singleton(newPartition);

            // All existing topics are not in the fetchedPartitions, so they should be marked as
            // removed
            Set<TopicPartition> expectedRemovedPartitions = new HashSet<>();
            for (int i = 0; i < KafkaSourceTestEnv.NUM_PARTITIONS; i++) {
                expectedRemovedPartitions.add(new TopicPartition(TOPIC1, i));
                expectedRemovedPartitions.add(new TopicPartition(TOPIC2, i));
            }

            assertThat(partitionChange.getNewPartitions()).isEqualTo(expectedNewPartitions);
            assertThat(partitionChange.getRemovedPartitions()).isEqualTo(expectedRemovedPartitions);
        }
    }

    // -------------- some common startup sequence ---------------

    private void startEnumeratorAndRegisterReaders(
            MockSplitEnumeratorContext<KafkaPartitionSplit> context,
            KafkaSourceEnumerator enumerator)
            throws Throwable {
        // Start the enumerator and it should schedule a one time task to discover and assign
        // partitions.
        enumerator.start();

        // register reader 0 before the partition discovery.
        registerReader(context, enumerator, READER0);
        assertThat(context.getSplitsAssignmentSequence()).isEmpty();

        // Run the partition discover callable and check the partition assignment.
        runPeriodicPartitionDiscovery(context);
        verifyLastReadersAssignments(
                context, Collections.singleton(READER0), PRE_EXISTING_TOPICS, 1);

        // Register reader 1 after first partition discovery.
        registerReader(context, enumerator, READER1);
        verifyLastReadersAssignments(
                context, Collections.singleton(READER1), PRE_EXISTING_TOPICS, 2);
    }

    // ----------------------------------------

    private KafkaSourceEnumerator createEnumerator(
            MockSplitEnumeratorContext<KafkaPartitionSplit> enumContext,
            boolean enablePeriodicPartitionDiscovery) {
        return createEnumerator(
                enumContext, enablePeriodicPartitionDiscovery, EXCLUDE_DYNAMIC_TOPIC);
    }

    private KafkaSourceEnumerator createEnumerator(
            MockSplitEnumeratorContext<KafkaPartitionSplit> enumContext,
            long partitionDiscoveryIntervalMs) {
        return createEnumerator(enumContext, partitionDiscoveryIntervalMs, EXCLUDE_DYNAMIC_TOPIC);
    }

    private KafkaSourceEnumerator createEnumerator(
            MockSplitEnumeratorContext<KafkaPartitionSplit> enumContext,
            boolean enablePeriodicPartitionDiscovery,
            boolean includeDynamicTopic) {
        List<String> topics = new ArrayList<>(PRE_EXISTING_TOPICS);
        if (includeDynamicTopic) {
            topics.add(DYNAMIC_TOPIC_NAME);
        }
        return createEnumerator(
                enumContext,
                enablePeriodicPartitionDiscovery ? 1 : -1,
                topics,
                Collections.emptySet(),
                new Properties());
    }

    private KafkaSourceEnumerator createEnumerator(
            MockSplitEnumeratorContext<KafkaPartitionSplit> enumContext,
            long partitionDiscoveryIntervalMs,
            boolean includeDynamicTopic) {
        List<String> topics = new ArrayList<>(PRE_EXISTING_TOPICS);
        if (includeDynamicTopic) {
            topics.add(DYNAMIC_TOPIC_NAME);
        }
        return createEnumerator(
                enumContext,
                partitionDiscoveryIntervalMs,
                topics,
                Collections.emptySet(),
                new Properties());
    }

    /**
     * Create the enumerator. For the purpose of the tests in this class we don't care about the
     * subscriber and offsets initializer, so just use arbitrary settings.
     */
    private KafkaSourceEnumerator createEnumerator(
            MockSplitEnumeratorContext<KafkaPartitionSplit> enumContext,
            long partitionDiscoveryIntervalMs,
            Collection<String> topicsToSubscribe,
            Set<TopicPartition> assignedPartitions,
            Properties overrideProperties) {
        // Use a TopicPatternSubscriber so that no exception if a subscribed topic hasn't been
        // created yet.
        StringJoiner topicNameJoiner = new StringJoiner("|");
        topicsToSubscribe.forEach(topicNameJoiner::add);
        Pattern topicPattern = Pattern.compile(topicNameJoiner.toString());
        KafkaSubscriber subscriber = KafkaSubscriber.getTopicPatternSubscriber(topicPattern);

        OffsetsInitializer startingOffsetsInitializer = OffsetsInitializer.earliest();
        OffsetsInitializer stoppingOffsetsInitializer = new NoStoppingOffsetsInitializer();

        Properties props =
                new Properties(KafkaSourceTestEnv.getConsumerProperties(StringDeserializer.class));
        KafkaSourceEnumerator.deepCopyProperties(overrideProperties, props);
        props.setProperty(
                KafkaSourceOptions.PARTITION_DISCOVERY_INTERVAL_MS.key(),
                String.valueOf(partitionDiscoveryIntervalMs));

        return new KafkaSourceEnumerator(
                subscriber,
                startingOffsetsInitializer,
                stoppingOffsetsInitializer,
                props,
                enumContext,
                Boundedness.CONTINUOUS_UNBOUNDED,
                assignedPartitions);
    }

    // ---------------------

    private void registerReader(
            MockSplitEnumeratorContext<KafkaPartitionSplit> context,
            KafkaSourceEnumerator enumerator,
            int reader) {
        context.registerReader(new ReaderInfo(reader, "location 0"));
        enumerator.addReader(reader);
    }

    private void verifyLastReadersAssignments(
            MockSplitEnumeratorContext<KafkaPartitionSplit> context,
            Collection<Integer> readers,
            Set<String> topics,
            int expectedAssignmentSeqSize) {
        verifyAssignments(
                getExpectedAssignments(new HashSet<>(readers), topics),
                context.getSplitsAssignmentSequence()
                        .get(expectedAssignmentSeqSize - 1)
                        .assignment());
    }

    private void verifyAssignments(
            Map<Integer, Set<TopicPartition>> expectedAssignments,
            Map<Integer, List<KafkaPartitionSplit>> actualAssignments) {
        actualAssignments.forEach(
                (reader, splits) -> {
                    Set<TopicPartition> expectedAssignmentsForReader =
                            expectedAssignments.get(reader);
                    assertThat(expectedAssignmentsForReader).isNotNull();
                    assertThat(splits.size()).isEqualTo(expectedAssignmentsForReader.size());
                    for (KafkaPartitionSplit split : splits) {
                        assertThat(expectedAssignmentsForReader)
                                .contains(split.getTopicPartition());
                    }
                });
    }

    private Map<Integer, Set<TopicPartition>> getExpectedAssignments(
            Set<Integer> readers, Set<String> topics) {
        Map<Integer, Set<TopicPartition>> expectedAssignments = new HashMap<>();
        Set<TopicPartition> allPartitions = new HashSet<>();

        if (topics.contains(DYNAMIC_TOPIC_NAME)) {
            for (int i = 0; i < NUM_PARTITIONS_DYNAMIC_TOPIC; i++) {
                allPartitions.add(new TopicPartition(DYNAMIC_TOPIC_NAME, i));
            }
        }

        for (TopicPartition tp : KafkaSourceTestEnv.getPartitionsForTopics(PRE_EXISTING_TOPICS)) {
            if (topics.contains(tp.topic())) {
                allPartitions.add(tp);
            }
        }

        for (TopicPartition tp : allPartitions) {
            int ownerReader = KafkaSourceEnumerator.getSplitOwner(tp, NUM_SUBTASKS);
            if (readers.contains(ownerReader)) {
                expectedAssignments.computeIfAbsent(ownerReader, r -> new HashSet<>()).add(tp);
            }
        }
        return expectedAssignments;
    }

    private void verifySplitAssignmentWithPartitions(
            Map<Integer, Set<TopicPartition>> expectedAssignment,
            Set<TopicPartition> actualTopicPartitions) {
        final Set<TopicPartition> allTopicPartitionsFromAssignment = new HashSet<>();
        expectedAssignment.forEach(
                (reader, topicPartitions) ->
                        allTopicPartitionsFromAssignment.addAll(topicPartitions));
        assertThat(actualTopicPartitions).isEqualTo(allTopicPartitionsFromAssignment);
    }

    private Set<TopicPartition> asEnumState(Map<Integer, List<KafkaPartitionSplit>> assignments) {
        Set<TopicPartition> enumState = new HashSet<>();
        assignments.forEach(
                (reader, assignment) ->
                        assignment.forEach(split -> enumState.add(split.getTopicPartition())));
        return enumState;
    }

    private void runOneTimePartitionDiscovery(
            MockSplitEnumeratorContext<KafkaPartitionSplit> context) throws Throwable {
        // Fetch potential topic descriptions
        context.runNextOneTimeCallable();
        // Initialize offsets for discovered partitions
        if (!context.getOneTimeCallables().isEmpty()) {
            context.runNextOneTimeCallable();
        }
    }

    private void runPeriodicPartitionDiscovery(
            MockSplitEnumeratorContext<KafkaPartitionSplit> context) throws Throwable {
        // Fetch potential topic descriptions
        context.runPeriodicCallable(PARTITION_DISCOVERY_CALLABLE_INDEX);
        // Initialize offsets for discovered partitions
        if (!context.getOneTimeCallables().isEmpty()) {
            context.runNextOneTimeCallable();
        }
    }
}