/*
 * 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.streaming.connectors.kinesis;

import org.apache.flink.api.common.ExecutionConfig;
import org.apache.flink.api.common.functions.RuntimeContext;
import org.apache.flink.api.common.serialization.DeserializationSchema;
import org.apache.flink.api.common.serialization.SimpleStringSchema;
import org.apache.flink.api.common.state.ListState;
import org.apache.flink.api.common.state.ListStateDescriptor;
import org.apache.flink.api.common.state.OperatorStateStore;
import org.apache.flink.api.common.time.Deadline;
import org.apache.flink.api.common.typeinfo.TypeInformation;
import org.apache.flink.api.common.typeutils.TypeSerializer;
import org.apache.flink.api.java.tuple.Tuple2;
import org.apache.flink.api.java.typeutils.TupleTypeInfo;
import org.apache.flink.api.java.typeutils.runtime.TupleSerializer;
import org.apache.flink.configuration.Configuration;
import org.apache.flink.mock.Whitebox;
import org.apache.flink.runtime.state.FunctionSnapshotContext;
import org.apache.flink.runtime.state.StateInitializationContext;
import org.apache.flink.runtime.state.StateSnapshotContextSynchronousImpl;
import org.apache.flink.streaming.api.TimeCharacteristic;
import org.apache.flink.streaming.api.functions.source.SourceFunction;
import org.apache.flink.streaming.api.functions.timestamps.BoundedOutOfOrdernessTimestampExtractor;
import org.apache.flink.streaming.api.operators.StreamSource;
import org.apache.flink.streaming.api.operators.collect.utils.MockFunctionSnapshotContext;
import org.apache.flink.streaming.api.watermark.Watermark;
import org.apache.flink.streaming.api.windowing.time.Time;
import org.apache.flink.streaming.connectors.kinesis.config.ConsumerConfigConstants;
import org.apache.flink.streaming.connectors.kinesis.internals.KinesisDataFetcher;
import org.apache.flink.streaming.connectors.kinesis.model.KinesisStreamShard;
import org.apache.flink.streaming.connectors.kinesis.model.KinesisStreamShardState;
import org.apache.flink.streaming.connectors.kinesis.model.SentinelSequenceNumber;
import org.apache.flink.streaming.connectors.kinesis.model.SequenceNumber;
import org.apache.flink.streaming.connectors.kinesis.model.StreamShardHandle;
import org.apache.flink.streaming.connectors.kinesis.model.StreamShardMetadata;
import org.apache.flink.streaming.connectors.kinesis.serialization.KinesisDeserializationSchema;
import org.apache.flink.streaming.connectors.kinesis.serialization.KinesisDeserializationSchemaWrapper;
import org.apache.flink.streaming.connectors.kinesis.testutils.FakeKinesisBehavioursFactory;
import org.apache.flink.streaming.connectors.kinesis.testutils.KinesisShardIdGenerator;
import org.apache.flink.streaming.connectors.kinesis.testutils.TestUtils;
import org.apache.flink.streaming.connectors.kinesis.testutils.TestableFlinkKinesisConsumer;
import org.apache.flink.streaming.connectors.kinesis.util.KinesisConfigUtil;
import org.apache.flink.streaming.connectors.kinesis.util.KinesisStateUtil;
import org.apache.flink.streaming.connectors.kinesis.util.RecordEmitter;
import org.apache.flink.streaming.connectors.kinesis.util.WatermarkTracker;
import org.apache.flink.streaming.util.AbstractStreamOperatorTestHarness;
import org.apache.flink.streaming.util.CollectingSourceContext;
import org.apache.flink.streaming.util.MockStreamingRuntimeContext;
import org.apache.flink.types.PojoTestUtils;
import org.apache.flink.util.InstantiationUtil;
import org.apache.flink.util.TestLogger;

import com.amazonaws.services.kinesis.model.HashKeyRange;
import com.amazonaws.services.kinesis.model.SequenceNumberRange;
import com.amazonaws.services.kinesis.model.Shard;
import org.junit.Test;
import org.junit.runner.RunWith;
import org.mockito.MockedStatic;
import org.mockito.Mockito;
import org.powermock.api.mockito.PowerMockito;
import org.powermock.core.classloader.annotations.PrepareForTest;
import org.powermock.modules.junit4.PowerMockRunner;

import java.io.Serializable;
import java.time.Duration;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collections;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Properties;
import java.util.UUID;
import java.util.concurrent.BlockingQueue;
import java.util.concurrent.ConcurrentLinkedQueue;
import java.util.concurrent.LinkedBlockingQueue;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.concurrent.atomic.AtomicLong;
import java.util.concurrent.atomic.AtomicReference;
import java.util.function.Supplier;

import static org.assertj.core.api.Assertions.assertThat;
import static org.mockito.ArgumentMatchers.any;
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;

/** Suite of FlinkKinesisConsumer tests for the methods called throughout the source life cycle. */
@RunWith(PowerMockRunner.class)
@PrepareForTest(FlinkKinesisConsumer.class)
public class FlinkKinesisConsumerTest extends TestLogger {

    // ----------------------------------------------------------------------
    // Tests related to state initialization
    // ----------------------------------------------------------------------

    @Test
    public void testUseRestoredStateForSnapshotIfFetcherNotInitialized() throws Exception {
        Properties config = TestUtils.getStandardProperties();

        List<Tuple2<StreamShardMetadata, SequenceNumber>> globalUnionState =
                Arrays.asList(
                        createShardState("fakeStream", 0, "1"),
                        createShardState("fakeStream", 1, "1"),
                        createShardState("fakeStream", 2, "1"),
                        createShardState("fakeStream", 3, "1"));

        TestingListState<Tuple2<StreamShardMetadata, SequenceNumber>> listState =
                new TestingListState<>();
        listState.addAll(globalUnionState);

        FlinkKinesisConsumer<String> consumer =
                new FlinkKinesisConsumer<>("fakeStream", new SimpleStringSchema(), config);
        RuntimeContext context = new MockStreamingRuntimeContext(true, 2, 0);
        consumer.setRuntimeContext(context);

        OperatorStateStore operatorStateStore = mock(OperatorStateStore.class);
        when(operatorStateStore.getUnionListState(any(ListStateDescriptor.class)))
                .thenReturn(listState);

        StateInitializationContext initializationContext = mock(StateInitializationContext.class);
        when(initializationContext.getOperatorStateStore()).thenReturn(operatorStateStore);
        when(initializationContext.isRestored()).thenReturn(true);

        consumer.initializeState(initializationContext);

        // only opened, not run
        consumer.open(new Configuration());

        // arbitrary checkpoint id and timestamp
        consumer.snapshotState(new StateSnapshotContextSynchronousImpl(123, 123));

        assertThat(listState.isClearCalled()).isTrue();

        // the checkpointed list state should contain only the shards that it should subscribe to
        assertThat(listState.getList()).hasSize(globalUnionState.size() / 2);
        assertThat(listState.getList()).contains(globalUnionState.get(0));
        assertThat(listState.getList()).contains(globalUnionState.get(2));
    }

    @Test
    public void testListStateChangedAfterSnapshotState() throws Exception {

        // ----------------------------------------------------------------------
        // setup config, initial state and expected state snapshot
        // ----------------------------------------------------------------------
        List<Tuple2<StreamShardMetadata, SequenceNumber>> initialState =
                Collections.singletonList(createShardState("fakeStream", 0, "1"));

        List<Tuple2<StreamShardMetadata, SequenceNumber>> expectedStateSnapshot =
                Arrays.asList(
                        createShardState("fakeStream", 0, "12"),
                        createShardState("fakeStream", 1, "11"),
                        createShardState("fakeStream", 2, "31"));

        // ----------------------------------------------------------------------
        // mock operator state backend and initial state for initializeState()
        // ----------------------------------------------------------------------

        TestingListState<Tuple2<StreamShardMetadata, SequenceNumber>> listState =
                new TestingListState<>();
        listState.addAll(initialState);

        // ----------------------------------------------------------------------
        // mock a running fetcher and its state for snapshot
        // ----------------------------------------------------------------------

        HashMap<StreamShardMetadata, SequenceNumber> stateSnapshot = new HashMap<>();
        for (Tuple2<StreamShardMetadata, SequenceNumber> tuple : expectedStateSnapshot) {
            stateSnapshot.put(tuple.f0, tuple.f1);
        }

        KinesisDataFetcher<String> mockedFetcher = mock(KinesisDataFetcher.class);
        when(mockedFetcher.snapshotState()).thenReturn(stateSnapshot);

        // ----------------------------------------------------------------------
        // create a consumer and test the snapshotState()
        // ----------------------------------------------------------------------

        FlinkKinesisConsumer<String> mockedConsumer =
                prepareMockedConsumer(
                        "fakeStream", new SimpleStringSchema(), mockedFetcher, listState);

        mockedConsumer.snapshotState(mock(FunctionSnapshotContext.class));

        // ----------------------------------------------------------------------
        // verify that state had been updated
        // ----------------------------------------------------------------------
        assertThat(listState.clearCalled).isTrue();
        assertThat(listState.getList())
                .hasSize(3)
                .doesNotContainAnyElementsOf(initialState)
                .containsAll(expectedStateSnapshot);
    }

    @Test
    public void testSnapshotStateChangesAfterCancel() throws Exception {

        // ----------------------------------------------------------------------
        // setup config, initial state and expected state snapshot
        // ----------------------------------------------------------------------
        List<Tuple2<StreamShardMetadata, SequenceNumber>> initialState =
                Collections.singletonList(createShardState("fakeStream", 0, "11"));

        List<Tuple2<StreamShardMetadata, SequenceNumber>> expectedStateSnapshot =
                Collections.singletonList(createShardState("fakeStream", 0, "12"));

        // ----------------------------------------------------------------------
        // mock operator state backend and initial state for initializeState()
        // ----------------------------------------------------------------------

        TestingListState<Tuple2<StreamShardMetadata, SequenceNumber>> listState =
                new TestingListState<>();
        listState.addAll(initialState);

        // ----------------------------------------------------------------------
        // mock a running fetcher and its state for snapshot
        // ----------------------------------------------------------------------

        HashMap<StreamShardMetadata, SequenceNumber> stateSnapshot = new HashMap<>();
        for (Tuple2<StreamShardMetadata, SequenceNumber> tuple : expectedStateSnapshot) {
            stateSnapshot.put(tuple.f0, tuple.f1);
        }

        KinesisDataFetcher<String> mockedFetcher = mock(KinesisDataFetcher.class);
        when(mockedFetcher.snapshotState()).thenReturn(stateSnapshot);

        // ----------------------------------------------------------------------
        // create a consumer and test the snapshotState()
        // ----------------------------------------------------------------------

        FlinkKinesisConsumer<String> mockedConsumer =
                prepareMockedConsumer(
                        "fakeStream", new SimpleStringSchema(), mockedFetcher, listState);

        mockedConsumer.cancel();
        mockedConsumer.snapshotState(new MockFunctionSnapshotContext(2));
        verify(mockedFetcher, times(1)).snapshotState();

        // ----------------------------------------------------------------------
        // verify that state had been updated
        // ----------------------------------------------------------------------
        assertThat(listState.isClearCalled()).isTrue();
        assertThat(listState.getList())
                .hasSize(1)
                .doesNotContainAnyElementsOf(initialState)
                .containsAll(expectedStateSnapshot);
    }

    @Test
    public void testSnapshotStateNotChangedAfterClose() throws Exception {

        // ----------------------------------------------------------------------
        // setup initial state
        // ----------------------------------------------------------------------

        List<Tuple2<StreamShardMetadata, SequenceNumber>> initialState =
                Collections.singletonList(createShardState("fakeStream", 0, "11"));

        // ----------------------------------------------------------------------
        // mock initial state
        // ----------------------------------------------------------------------

        TestingListState<Tuple2<StreamShardMetadata, SequenceNumber>> listState =
                new TestingListState<>();
        listState.addAll(initialState);

        // ----------------------------------------------------------------------
        // mock a running fetcher and its state for snapshot
        // ----------------------------------------------------------------------

        KinesisDataFetcher<String> mockedFetcher = mock(KinesisDataFetcher.class);
        when(mockedFetcher.snapshotState()).thenReturn(new HashMap<>());

        // ----------------------------------------------------------------------
        // create a consumer and test the snapshotState()
        // ----------------------------------------------------------------------

        FlinkKinesisConsumer<String> mockedConsumer =
                prepareMockedConsumer(
                        "fakeStream", new SimpleStringSchema(), mockedFetcher, listState);

        mockedConsumer.close();
        mockedConsumer.snapshotState(new MockFunctionSnapshotContext(3));
        verify(mockedFetcher, never()).snapshotState();
        assertThat(listState.isClearCalled()).isFalse();
        assertThat(listState.getList()).containsAll(initialState);
    }

    private <T> FlinkKinesisConsumer<T> prepareMockedConsumer(
            String streamName,
            DeserializationSchema<T> schema,
            KinesisDataFetcher<T> fetcher,
            ListState<?> listState)
            throws Exception {

        Properties config = TestUtils.getStandardProperties();

        OperatorStateStore operatorStateStore = mock(OperatorStateStore.class);
        when(operatorStateStore.getUnionListState(any(ListStateDescriptor.class)))
                .thenReturn(listState);

        StateInitializationContext initializationContext = mock(StateInitializationContext.class);
        when(initializationContext.getOperatorStateStore()).thenReturn(operatorStateStore);
        when(initializationContext.isRestored()).thenReturn(true);

        FlinkKinesisConsumer<T> consumer = new FlinkKinesisConsumer<>(streamName, schema, config);
        FlinkKinesisConsumer<T> mockedConsumer = spy(consumer);

        RuntimeContext context = new MockStreamingRuntimeContext(true, 1, 0);

        mockedConsumer.setRuntimeContext(context);
        mockedConsumer.initializeState(initializationContext);
        mockedConsumer.open(new Configuration());
        Whitebox.setInternalState(mockedConsumer, "fetcher", fetcher); // mock consumer as running.

        return mockedConsumer;
    }

    /**
     * Before using an explicit TypeSerializer for the state the {@link FlinkKinesisConsumer} was
     * creating a serializer implicitly using a {@link TypeInformation}. After fixing issue
     * FLINK-24943, * serializer is created explicitly. Here, we verify that previous approach is
     * compatible with the new one.
     */
    @Test
    public void testExplicitStateSerializerCompatibility() throws Exception {
        ExecutionConfig executionConfig = new ExecutionConfig();

        Tuple2<StreamShardMetadata, SequenceNumber> tuple = createShardState("fakeStream", 0, "1");

        // This is how serializer was created implicitly using a TypeInformation
        // and since SequenceNumber is GenericType, Flink falls back to Kryo
        TypeInformation<Tuple2<StreamShardMetadata, SequenceNumber>> originalShardsStateTypeInfo =
                new TupleTypeInfo<>(
                        TypeInformation.of(StreamShardMetadata.class),
                        TypeInformation.of(SequenceNumber.class));
        TypeSerializer<Tuple2<StreamShardMetadata, SequenceNumber>> serializerFromTypeInfo =
                originalShardsStateTypeInfo.createSerializer(executionConfig);
        byte[] bytes = InstantiationUtil.serializeToByteArray(serializerFromTypeInfo, tuple);

        // This is how we create serializer explicitly with Kryo
        TupleSerializer<Tuple2<StreamShardMetadata, SequenceNumber>> serializerFromKryo =
                KinesisStateUtil.createShardsStateSerializer(executionConfig);

        Tuple2<StreamShardMetadata, SequenceNumber> actualTuple =
                InstantiationUtil.deserializeFromByteArray(serializerFromKryo, bytes);

        // Both ways should be the same
        assertThat(tuple)
                .overridingErrorMessage(
                        "Explicit serializer is not compatible with "
                                + "implicit method of creating serializer using TypeInformation.")
                .isEqualTo(actualTuple);
    }

    private Tuple2<StreamShardMetadata, SequenceNumber> createShardState(
            String streamName, int shardNumber, String sequenceNumber) {
        return Tuple2.of(
                KinesisDataFetcher.convertToStreamShardMetadata(
                        new StreamShardHandle(
                                streamName,
                                new Shard()
                                        .withShardId(
                                                KinesisShardIdGenerator.generateFromShardOrder(
                                                        shardNumber)))),
                new SequenceNumber(sequenceNumber));
    }

    // ----------------------------------------------------------------------
    // Tests related to fetcher initialization
    // ----------------------------------------------------------------------

    @Test
    @SuppressWarnings("unchecked")
    public void testFetcherShouldNotBeRestoringFromFailureIfNotRestoringFromCheckpoint()
            throws Exception {
        KinesisDataFetcher mockedFetcher = mockKinesisDataFetcher();

        // assume the given config is correct
        try (MockedStatic<KinesisConfigUtil> kcu = mockStatic(KinesisConfigUtil.class)) {
            TestableFlinkKinesisConsumer consumer =
                    new TestableFlinkKinesisConsumer("fakeStream", new Properties(), 10, 2);
            consumer.open(new Configuration());
            consumer.run(Mockito.mock(SourceFunction.SourceContext.class));
        }
    }

    @Test
    @SuppressWarnings("unchecked")
    public void testFetcherShouldBeCorrectlySeededIfRestoringFromCheckpoint() throws Exception {

        // ----------------------------------------------------------------------
        // setup initial state
        // ----------------------------------------------------------------------

        HashMap<StreamShardHandle, SequenceNumber> fakeRestoredState = getFakeRestoredStore("all");

        // ----------------------------------------------------------------------
        // mock operator state backend and initial state for initializeState()
        // ----------------------------------------------------------------------

        TestingListState<Tuple2<StreamShardMetadata, SequenceNumber>> listState =
                new TestingListState<>();
        for (Map.Entry<StreamShardHandle, SequenceNumber> state : fakeRestoredState.entrySet()) {
            listState.add(
                    Tuple2.of(
                            KinesisDataFetcher.convertToStreamShardMetadata(state.getKey()),
                            state.getValue()));
        }

        OperatorStateStore operatorStateStore = mock(OperatorStateStore.class);
        when(operatorStateStore.getUnionListState(any(ListStateDescriptor.class)))
                .thenReturn(listState);

        StateInitializationContext initializationContext = mock(StateInitializationContext.class);
        when(initializationContext.getOperatorStateStore()).thenReturn(operatorStateStore);
        when(initializationContext.isRestored()).thenReturn(true);

        // ----------------------------------------------------------------------
        // mock fetcher
        // ----------------------------------------------------------------------

        KinesisDataFetcher mockedFetcher = mockKinesisDataFetcher();
        List<StreamShardHandle> shards = new ArrayList<>();
        shards.addAll(fakeRestoredState.keySet());
        when(mockedFetcher.discoverNewShardsToSubscribe()).thenReturn(shards);

        // assume the given config is correct
        try (MockedStatic<KinesisConfigUtil> kcu = mockStatic(KinesisConfigUtil.class)) {

            // ----------------------------------------------------------------------
            // start to test fetcher's initial state seeding
            // ----------------------------------------------------------------------

            TestableFlinkKinesisConsumer consumer =
                    new TestableFlinkKinesisConsumer("fakeStream", new Properties(), 10, 2);
            consumer.initializeState(initializationContext);
            consumer.open(new Configuration());
            consumer.run(Mockito.mock(SourceFunction.SourceContext.class));

            for (Map.Entry<StreamShardHandle, SequenceNumber> restoredShard :
                    fakeRestoredState.entrySet()) {
                Mockito.verify(mockedFetcher)
                        .registerNewSubscribedShardState(
                                new KinesisStreamShardState(
                                        KinesisDataFetcher.convertToStreamShardMetadata(
                                                restoredShard.getKey()),
                                        restoredShard.getKey(),
                                        restoredShard.getValue()));
            }
        }
    }

    @Test
    @SuppressWarnings("unchecked")
    public void testFetcherShouldBeCorrectlySeededOnlyItsOwnStates() throws Exception {

        // ----------------------------------------------------------------------
        // setup initial state
        // ----------------------------------------------------------------------

        HashMap<StreamShardHandle, SequenceNumber> fakeRestoredState =
                getFakeRestoredStore("fakeStream1");

        HashMap<StreamShardHandle, SequenceNumber> fakeRestoredStateForOthers =
                getFakeRestoredStore("fakeStream2");

        // ----------------------------------------------------------------------
        // mock operator state backend and initial state for initializeState()
        // ----------------------------------------------------------------------

        TestingListState<Tuple2<StreamShardMetadata, SequenceNumber>> listState =
                new TestingListState<>();
        for (Map.Entry<StreamShardHandle, SequenceNumber> state : fakeRestoredState.entrySet()) {
            listState.add(
                    Tuple2.of(
                            KinesisDataFetcher.convertToStreamShardMetadata(state.getKey()),
                            state.getValue()));
        }
        for (Map.Entry<StreamShardHandle, SequenceNumber> state :
                fakeRestoredStateForOthers.entrySet()) {
            listState.add(
                    Tuple2.of(
                            KinesisDataFetcher.convertToStreamShardMetadata(state.getKey()),
                            state.getValue()));
        }

        OperatorStateStore operatorStateStore = mock(OperatorStateStore.class);
        when(operatorStateStore.getUnionListState(any(ListStateDescriptor.class)))
                .thenReturn(listState);

        StateInitializationContext initializationContext = mock(StateInitializationContext.class);
        when(initializationContext.getOperatorStateStore()).thenReturn(operatorStateStore);
        when(initializationContext.isRestored()).thenReturn(true);

        // ----------------------------------------------------------------------
        // mock fetcher
        // ----------------------------------------------------------------------

        KinesisDataFetcher mockedFetcher = mockKinesisDataFetcher();
        List<StreamShardHandle> shards = new ArrayList<>();
        shards.addAll(fakeRestoredState.keySet());
        when(mockedFetcher.discoverNewShardsToSubscribe()).thenReturn(shards);

        // assume the given config is correct
        try (MockedStatic<KinesisConfigUtil> kcu = mockStatic(KinesisConfigUtil.class)) {

            // ----------------------------------------------------------------------
            // start to test fetcher's initial state seeding
            // ----------------------------------------------------------------------

            TestableFlinkKinesisConsumer consumer =
                    new TestableFlinkKinesisConsumer("fakeStream", new Properties(), 10, 2);
            consumer.initializeState(initializationContext);
            consumer.open(new Configuration());
            consumer.run(Mockito.mock(SourceFunction.SourceContext.class));

            for (Map.Entry<StreamShardHandle, SequenceNumber> restoredShard :
                    fakeRestoredStateForOthers.entrySet()) {
                // should never get restored state not belonging to itself
                Mockito.verify(mockedFetcher, never())
                        .registerNewSubscribedShardState(
                                new KinesisStreamShardState(
                                        KinesisDataFetcher.convertToStreamShardMetadata(
                                                restoredShard.getKey()),
                                        restoredShard.getKey(),
                                        restoredShard.getValue()));
            }
            for (Map.Entry<StreamShardHandle, SequenceNumber> restoredShard :
                    fakeRestoredState.entrySet()) {
                // should get restored state belonging to itself
                Mockito.verify(mockedFetcher)
                        .registerNewSubscribedShardState(
                                new KinesisStreamShardState(
                                        KinesisDataFetcher.convertToStreamShardMetadata(
                                                restoredShard.getKey()),
                                        restoredShard.getKey(),
                                        restoredShard.getValue()));
            }
        }
    }

    /*
     * This tests that the consumer correctly picks up shards that were not discovered on the previous run.
     *
     * Case under test:
     *
     * If the original parallelism is 2 and states are:
     *   Consumer subtask 1:
     *     stream1, shard1, SequentialNumber(xxx)
     *   Consumer subtask 2:
     *     stream1, shard2, SequentialNumber(yyy)
     *
     * After discoverNewShardsToSubscribe() if there were two shards (shard3, shard4) created:
     *   Consumer subtask 1 (late for discoverNewShardsToSubscribe()):
     *     stream1, shard1, SequentialNumber(xxx)
     *   Consumer subtask 2:
     *     stream1, shard2, SequentialNumber(yyy)
     *     stream1, shard4, SequentialNumber(zzz)
     *
     * If snapshotState() occurs and parallelism is changed to 1:
     *   Union state will be:
     *     stream1, shard1, SequentialNumber(xxx)
     *     stream1, shard2, SequentialNumber(yyy)
     *     stream1, shard4, SequentialNumber(zzz)
     *   Fetcher should be seeded with:
     *     stream1, shard1, SequentialNumber(xxx)
     *     stream1, shard2, SequentialNumber(yyy)
     *     stream1, share3, SentinelSequenceNumber.SENTINEL_EARLIEST_SEQUENCE_NUM
     *     stream1, shard4, SequentialNumber(zzz)
     */
    @Test
    @SuppressWarnings("unchecked")
    public void testFetcherShouldBeCorrectlySeededWithNewDiscoveredKinesisStreamShard()
            throws Exception {

        // ----------------------------------------------------------------------
        // setup initial state
        // ----------------------------------------------------------------------

        HashMap<StreamShardHandle, SequenceNumber> fakeRestoredState = getFakeRestoredStore("all");

        // ----------------------------------------------------------------------
        // mock operator state backend and initial state for initializeState()
        // ----------------------------------------------------------------------

        TestingListState<Tuple2<StreamShardMetadata, SequenceNumber>> listState =
                new TestingListState<>();
        for (Map.Entry<StreamShardHandle, SequenceNumber> state : fakeRestoredState.entrySet()) {
            listState.add(
                    Tuple2.of(
                            KinesisDataFetcher.convertToStreamShardMetadata(state.getKey()),
                            state.getValue()));
        }

        OperatorStateStore operatorStateStore = mock(OperatorStateStore.class);
        when(operatorStateStore.getUnionListState(any(ListStateDescriptor.class)))
                .thenReturn(listState);

        StateInitializationContext initializationContext = mock(StateInitializationContext.class);
        when(initializationContext.getOperatorStateStore()).thenReturn(operatorStateStore);
        when(initializationContext.isRestored()).thenReturn(true);

        // ----------------------------------------------------------------------
        // mock fetcher
        // ----------------------------------------------------------------------

        KinesisDataFetcher mockedFetcher = mockKinesisDataFetcher();
        List<StreamShardHandle> shards = new ArrayList<>();
        shards.addAll(fakeRestoredState.keySet());
        shards.add(
                new StreamShardHandle(
                        "fakeStream2",
                        new Shard()
                                .withShardId(KinesisShardIdGenerator.generateFromShardOrder(2))));
        when(mockedFetcher.discoverNewShardsToSubscribe()).thenReturn(shards);

        // assume the given config is correct
        try (MockedStatic<KinesisConfigUtil> kcu = mockStatic(KinesisConfigUtil.class)) {

            // ----------------------------------------------------------------------
            // start to test fetcher's initial state seeding
            // ----------------------------------------------------------------------

            TestableFlinkKinesisConsumer consumer =
                    new TestableFlinkKinesisConsumer("fakeStream", new Properties(), 10, 2);
            consumer.initializeState(initializationContext);
            consumer.open(new Configuration());
            consumer.run(Mockito.mock(SourceFunction.SourceContext.class));

            fakeRestoredState.put(
                    new StreamShardHandle(
                            "fakeStream2",
                            new Shard()
                                    .withShardId(
                                            KinesisShardIdGenerator.generateFromShardOrder(2))),
                    SentinelSequenceNumber.SENTINEL_EARLIEST_SEQUENCE_NUM.get());
            for (Map.Entry<StreamShardHandle, SequenceNumber> restoredShard :
                    fakeRestoredState.entrySet()) {
                Mockito.verify(mockedFetcher)
                        .registerNewSubscribedShardState(
                                new KinesisStreamShardState(
                                        KinesisDataFetcher.convertToStreamShardMetadata(
                                                restoredShard.getKey()),
                                        restoredShard.getKey(),
                                        restoredShard.getValue()));
            }
        }
    }

    @Test
    public void testLegacyKinesisStreamShardToStreamShardMetadataConversion() {
        String streamName = "fakeStream1";
        String shardId = "shard-000001";
        String parentShardId = "shard-000002";
        String adjacentParentShardId = "shard-000003";
        String startingHashKey = "key-000001";
        String endingHashKey = "key-000010";
        String startingSequenceNumber = "seq-0000021";
        String endingSequenceNumber = "seq-00000031";

        StreamShardMetadata streamShardMetadata = new StreamShardMetadata();
        streamShardMetadata.setStreamName(streamName);
        streamShardMetadata.setShardId(shardId);
        streamShardMetadata.setParentShardId(parentShardId);
        streamShardMetadata.setAdjacentParentShardId(adjacentParentShardId);
        streamShardMetadata.setStartingHashKey(startingHashKey);
        streamShardMetadata.setEndingHashKey(endingHashKey);
        streamShardMetadata.setStartingSequenceNumber(startingSequenceNumber);
        streamShardMetadata.setEndingSequenceNumber(endingSequenceNumber);

        Shard shard =
                new Shard()
                        .withShardId(shardId)
                        .withParentShardId(parentShardId)
                        .withAdjacentParentShardId(adjacentParentShardId)
                        .withHashKeyRange(
                                new HashKeyRange()
                                        .withStartingHashKey(startingHashKey)
                                        .withEndingHashKey(endingHashKey))
                        .withSequenceNumberRange(
                                new SequenceNumberRange()
                                        .withStartingSequenceNumber(startingSequenceNumber)
                                        .withEndingSequenceNumber(endingSequenceNumber));
        KinesisStreamShard kinesisStreamShard = new KinesisStreamShard(streamName, shard);

        assertThat(KinesisStreamShard.convertToStreamShardMetadata(kinesisStreamShard))
                .isEqualTo(streamShardMetadata);
    }

    @Test
    public void testStreamShardMetadataSerializedUsingPojoSerializer() {
        PojoTestUtils.assertSerializedAsPojo(StreamShardMetadata.class);
    }

    /**
     * FLINK-8484: ensure that a state change in the StreamShardMetadata other than {@link
     * StreamShardMetadata#getShardId()} or {@link StreamShardMetadata#getStreamName()} does not
     * result in the shard not being able to be restored. This handles the corner case where the
     * stored shard metadata is open (no ending sequence number), but after the job restore, the
     * shard has been closed (ending number set) due to re-sharding, and we can no longer rely on
     * {@link StreamShardMetadata#equals(Object)} to find back the sequence number in the collection
     * of restored shard metadata.
     *
     * <p>Therefore, we will rely on synchronizing the snapshot's state with the Kinesis shard
     * before attempting to find back the sequence number to restore.
     */
    @Test
    public void testFindSequenceNumberToRestoreFromIfTheShardHasBeenClosedSinceTheStateWasStored()
            throws Exception {
        // ----------------------------------------------------------------------
        // setup initial state
        // ----------------------------------------------------------------------

        HashMap<StreamShardHandle, SequenceNumber> fakeRestoredState = getFakeRestoredStore("all");

        // ----------------------------------------------------------------------
        // mock operator state backend and initial state for initializeState()
        // ----------------------------------------------------------------------

        TestingListState<Tuple2<StreamShardMetadata, SequenceNumber>> listState =
                new TestingListState<>();
        for (Map.Entry<StreamShardHandle, SequenceNumber> state : fakeRestoredState.entrySet()) {
            listState.add(
                    Tuple2.of(
                            KinesisDataFetcher.convertToStreamShardMetadata(state.getKey()),
                            state.getValue()));
        }

        OperatorStateStore operatorStateStore = mock(OperatorStateStore.class);
        when(operatorStateStore.getUnionListState(any(ListStateDescriptor.class)))
                .thenReturn(listState);

        StateInitializationContext initializationContext = mock(StateInitializationContext.class);
        when(initializationContext.getOperatorStateStore()).thenReturn(operatorStateStore);
        when(initializationContext.isRestored()).thenReturn(true);

        // ----------------------------------------------------------------------
        // mock fetcher
        // ----------------------------------------------------------------------

        KinesisDataFetcher mockedFetcher = mockKinesisDataFetcher();
        List<StreamShardHandle> shards = new ArrayList<>();

        // create a fake stream shard handle based on the first entry in the restored state
        final StreamShardHandle originalStreamShardHandle =
                fakeRestoredState.keySet().iterator().next();
        final StreamShardHandle closedStreamShardHandle =
                new StreamShardHandle(
                        originalStreamShardHandle.getStreamName(),
                        originalStreamShardHandle.getShard());
        // close the shard handle by setting an ending sequence number
        final SequenceNumberRange sequenceNumberRange = new SequenceNumberRange();
        sequenceNumberRange.setEndingSequenceNumber("1293844");
        closedStreamShardHandle.getShard().setSequenceNumberRange(sequenceNumberRange);

        shards.add(closedStreamShardHandle);

        when(mockedFetcher.discoverNewShardsToSubscribe()).thenReturn(shards);

        // assume the given config is correct
        try (MockedStatic<KinesisConfigUtil> kcu = mockStatic(KinesisConfigUtil.class)) {

            // ----------------------------------------------------------------------
            // start to test fetcher's initial state seeding
            // ----------------------------------------------------------------------

            TestableFlinkKinesisConsumer consumer =
                    new TestableFlinkKinesisConsumer("fakeStream", new Properties(), 10, 2);
            consumer.initializeState(initializationContext);
            consumer.open(new Configuration());
            consumer.run(Mockito.mock(SourceFunction.SourceContext.class));

            Mockito.verify(mockedFetcher)
                    .registerNewSubscribedShardState(
                            new KinesisStreamShardState(
                                    KinesisDataFetcher.convertToStreamShardMetadata(
                                            closedStreamShardHandle),
                                    closedStreamShardHandle,
                                    fakeRestoredState.get(closedStreamShardHandle)));
        }
    }

    private static final class TestingListState<T> implements ListState<T> {

        private final List<T> list = new ArrayList<>();
        private boolean clearCalled = false;

        @Override
        public void clear() {
            list.clear();
            clearCalled = true;
        }

        @Override
        public Iterable<T> get() throws Exception {
            return list;
        }

        @Override
        public void add(T value) throws Exception {
            list.add(value);
        }

        public List<T> getList() {
            return list;
        }

        public boolean isClearCalled() {
            return clearCalled;
        }

        @Override
        public void update(List<T> values) throws Exception {
            list.clear();

            addAll(values);
        }

        @Override
        public void addAll(List<T> values) throws Exception {
            if (values != null) {
                list.addAll(values);
            }
        }
    }

    private HashMap<StreamShardHandle, SequenceNumber> getFakeRestoredStore(String streamName) {
        HashMap<StreamShardHandle, SequenceNumber> fakeRestoredState = new HashMap<>();

        if (streamName.equals("fakeStream1") || streamName.equals("all")) {
            fakeRestoredState.put(
                    new StreamShardHandle(
                            "fakeStream1",
                            new Shard()
                                    .withShardId(
                                            KinesisShardIdGenerator.generateFromShardOrder(0))),
                    new SequenceNumber(UUID.randomUUID().toString()));
            fakeRestoredState.put(
                    new StreamShardHandle(
                            "fakeStream1",
                            new Shard()
                                    .withShardId(
                                            KinesisShardIdGenerator.generateFromShardOrder(1))),
                    new SequenceNumber(UUID.randomUUID().toString()));
            fakeRestoredState.put(
                    new StreamShardHandle(
                            "fakeStream1",
                            new Shard()
                                    .withShardId(
                                            KinesisShardIdGenerator.generateFromShardOrder(2))),
                    new SequenceNumber(UUID.randomUUID().toString()));
        }

        if (streamName.equals("fakeStream2") || streamName.equals("all")) {
            fakeRestoredState.put(
                    new StreamShardHandle(
                            "fakeStream2",
                            new Shard()
                                    .withShardId(
                                            KinesisShardIdGenerator.generateFromShardOrder(0))),
                    new SequenceNumber(UUID.randomUUID().toString()));
            fakeRestoredState.put(
                    new StreamShardHandle(
                            "fakeStream2",
                            new Shard()
                                    .withShardId(
                                            KinesisShardIdGenerator.generateFromShardOrder(1))),
                    new SequenceNumber(UUID.randomUUID().toString()));
        }

        return fakeRestoredState;
    }

    private static KinesisDataFetcher mockKinesisDataFetcher() throws Exception {
        KinesisDataFetcher mockedFetcher = Mockito.mock(KinesisDataFetcher.class);

        java.lang.reflect.Constructor<KinesisDataFetcher> ctor =
                (java.lang.reflect.Constructor<KinesisDataFetcher>)
                        KinesisDataFetcher.class.getConstructors()[0];
        Class<?>[] otherParamTypes = new Class<?>[ctor.getParameterCount() - 1];
        System.arraycopy(
                ctor.getParameterTypes(), 1, otherParamTypes, 0, ctor.getParameterCount() - 1);

        Supplier<Object[]> argumentSupplier =
                () -> {
                    Object[] otherParamArgs = new Object[otherParamTypes.length];
                    for (int i = 0; i < otherParamTypes.length; i++) {
                        otherParamArgs[i] = Mockito.nullable(otherParamTypes[i]);
                    }
                    return otherParamArgs;
                };
        PowerMockito.whenNew(ctor)
                .withArguments(Mockito.any(ctor.getParameterTypes()[0]), argumentSupplier.get())
                .thenReturn(mockedFetcher);
        return mockedFetcher;
    }

    @Test
    public void testPeriodicWatermark() throws Exception {
        String streamName = "fakeStreamName";
        Time maxOutOfOrderness = Time.milliseconds(5);
        long autoWatermarkInterval = 1_000;

        HashMap<String, String> subscribedStreamsToLastDiscoveredShardIds = new HashMap<>();
        subscribedStreamsToLastDiscoveredShardIds.put(streamName, null);

        KinesisDeserializationSchema<String> deserializationSchema =
                new KinesisDeserializationSchemaWrapper<>(new SimpleStringSchema());
        Properties props = new Properties();
        props.setProperty(ConsumerConfigConstants.AWS_REGION, "us-east-1");
        props.setProperty(
                ConsumerConfigConstants.SHARD_GETRECORDS_INTERVAL_MILLIS, Long.toString(10L));

        BlockingQueue<String> shard1 = new LinkedBlockingQueue<>();
        BlockingQueue<String> shard2 = new LinkedBlockingQueue<>();

        Map<String, List<BlockingQueue<String>>> streamToQueueMap = new HashMap<>();
        streamToQueueMap.put(streamName, Arrays.asList(shard1, shard2));

        // override createFetcher to mock Kinesis
        FlinkKinesisConsumer<String> sourceFunc =
                new FlinkKinesisConsumer<String>(streamName, deserializationSchema, props) {
                    @Override
                    protected KinesisDataFetcher<String> createFetcher(
                            List<String> streams,
                            SourceContext<String> sourceContext,
                            RuntimeContext runtimeContext,
                            Properties configProps,
                            KinesisDeserializationSchema<String> deserializationSchema) {

                        KinesisDataFetcher<String> fetcher =
                                new KinesisDataFetcher<String>(
                                        streams,
                                        sourceContext,
                                        sourceContext.getCheckpointLock(),
                                        runtimeContext,
                                        configProps,
                                        deserializationSchema,
                                        getShardAssigner(),
                                        getPeriodicWatermarkAssigner(),
                                        null,
                                        new AtomicReference<>(),
                                        new ArrayList<>(),
                                        subscribedStreamsToLastDiscoveredShardIds,
                                        (props) ->
                                                FakeKinesisBehavioursFactory
                                                        .blockingQueueGetRecords(streamToQueueMap),
                                        null) {};
                        return fetcher;
                    }
                };

        sourceFunc.setShardAssigner(
                (streamShardHandle, i) -> {
                    // shardId-000000000000
                    return Integer.parseInt(
                            streamShardHandle
                                    .getShard()
                                    .getShardId()
                                    .substring("shardId-".length()));
                });

        sourceFunc.setPeriodicWatermarkAssigner(new TestTimestampExtractor(maxOutOfOrderness));

        // there is currently no test harness specifically for sources,
        // so we overlay the source thread here
        AbstractStreamOperatorTestHarness<Object> testHarness =
                new AbstractStreamOperatorTestHarness<Object>(
                        new StreamSource(sourceFunc), 1, 1, 0);
        testHarness.setTimeCharacteristic(TimeCharacteristic.EventTime);
        testHarness.getExecutionConfig().setAutoWatermarkInterval(autoWatermarkInterval);

        testHarness.initializeEmptyState();
        testHarness.open();

        ConcurrentLinkedQueue<Watermark> watermarks = new ConcurrentLinkedQueue<>();

        @SuppressWarnings("unchecked")
        SourceFunction.SourceContext<String> sourceContext =
                new CollectingSourceContext(
                        testHarness.getCheckpointLock(), testHarness.getOutput()) {
                    @Override
                    public void emitWatermark(Watermark mark) {
                        watermarks.add(mark);
                    }

                    @Override
                    public void markAsTemporarilyIdle() {}
                };

        new Thread(
                        () -> {
                            try {
                                sourceFunc.run(sourceContext);
                            } catch (InterruptedException e) {
                                // expected on cancel
                            } catch (Exception e) {
                                throw new RuntimeException(e);
                            }
                        })
                .start();

        shard1.put("1");
        shard1.put("2");
        shard2.put("10");
        int recordCount = 3;
        int watermarkCount = 0;
        awaitRecordCount(testHarness.getOutput(), recordCount);

        // Trigger watermark emit, first watermark is -3
        // - Shard-1 @2
        // - Shard-2 @10
        // - Watermark = min(2, 10) - maxOutOfOrderness = 2 - 5 = -3
        testHarness.setProcessingTime(testHarness.getProcessingTime() + autoWatermarkInterval);
        watermarkCount++;

        // advance watermark
        shard1.put("10");
        recordCount++;
        awaitRecordCount(testHarness.getOutput(), recordCount);

        // Trigger watermark emit, second watermark is -3
        // - Shard-1 @10
        // - Shard-2 @10
        // - Watermark = min(10, 10) - maxOutOfOrderness = 10 - 5 = 5
        testHarness.setProcessingTime(testHarness.getProcessingTime() + autoWatermarkInterval);
        watermarkCount++;

        sourceFunc.cancel();
        testHarness.close();

        assertThat(testHarness.getOutput()).as("record count").hasSize(recordCount);
        assertThat(watermarks).contains(new Watermark(-3), new Watermark(5));
        assertThat(watermarks).as("watermark count").hasSize(watermarkCount);
    }

    @Test
    public void testSourceSynchronization() throws Exception {

        final String streamName = "fakeStreamName";
        final Time maxOutOfOrderness = Time.milliseconds(5);
        final long autoWatermarkInterval = 1_000;
        final long watermarkSyncInterval = autoWatermarkInterval + 1;

        TestWatermarkTracker.WATERMARK.set(0);
        HashMap<String, String> subscribedStreamsToLastDiscoveredShardIds = new HashMap<>();
        subscribedStreamsToLastDiscoveredShardIds.put(streamName, null);

        final KinesisDeserializationSchema<String> deserializationSchema =
                new KinesisDeserializationSchemaWrapper<>(new OpenCheckingStringSchema());
        Properties props = new Properties();
        props.setProperty(ConsumerConfigConstants.AWS_REGION, "us-east-1");
        props.setProperty(
                ConsumerConfigConstants.SHARD_GETRECORDS_INTERVAL_MILLIS, Long.toString(10L));
        props.setProperty(
                ConsumerConfigConstants.WATERMARK_SYNC_MILLIS,
                Long.toString(watermarkSyncInterval));
        props.setProperty(ConsumerConfigConstants.WATERMARK_LOOKAHEAD_MILLIS, Long.toString(5));

        BlockingQueue<String> shard1 = new LinkedBlockingQueue<>();

        Map<String, List<BlockingQueue<String>>> streamToQueueMap = new HashMap<>();
        streamToQueueMap.put(streamName, Collections.singletonList(shard1));

        // override createFetcher to mock Kinesis
        FlinkKinesisConsumer<String> sourceFunc =
                new FlinkKinesisConsumer<String>(streamName, deserializationSchema, props) {
                    @Override
                    protected KinesisDataFetcher<String> createFetcher(
                            List<String> streams,
                            SourceFunction.SourceContext<String> sourceContext,
                            RuntimeContext runtimeContext,
                            Properties configProps,
                            KinesisDeserializationSchema<String> deserializationSchema) {

                        KinesisDataFetcher<String> fetcher =
                                new KinesisDataFetcher<String>(
                                        streams,
                                        sourceContext,
                                        sourceContext.getCheckpointLock(),
                                        runtimeContext,
                                        configProps,
                                        deserializationSchema,
                                        getShardAssigner(),
                                        getPeriodicWatermarkAssigner(),
                                        getWatermarkTracker(),
                                        new AtomicReference<>(),
                                        new ArrayList<>(),
                                        subscribedStreamsToLastDiscoveredShardIds,
                                        (props) ->
                                                FakeKinesisBehavioursFactory
                                                        .blockingQueueGetRecords(streamToQueueMap),
                                        null) {
                                    @Override
                                    protected void emitWatermark() {
                                        // necessary in this test to ensure that watermark state is
                                        // updated
                                        // before the watermark timer callback is triggered
                                        synchronized (sourceContext.getCheckpointLock()) {
                                            super.emitWatermark();
                                        }
                                    }
                                };
                        return fetcher;
                    }
                };

        sourceFunc.setShardAssigner(
                (streamShardHandle, i) -> {
                    // shardId-000000000000
                    return Integer.parseInt(
                            streamShardHandle
                                    .getShard()
                                    .getShardId()
                                    .substring("shardId-".length()));
                });

        sourceFunc.setPeriodicWatermarkAssigner(new TestTimestampExtractor(maxOutOfOrderness));

        sourceFunc.setWatermarkTracker(new TestWatermarkTracker());

        // there is currently no test harness specifically for sources,
        // so we overlay the source thread here
        AbstractStreamOperatorTestHarness<Object> testHarness =
                new AbstractStreamOperatorTestHarness<Object>(
                        new StreamSource(sourceFunc), 1, 1, 0);
        testHarness.setTimeCharacteristic(TimeCharacteristic.EventTime);
        testHarness.getExecutionConfig().setAutoWatermarkInterval(autoWatermarkInterval);

        testHarness.initializeEmptyState();
        testHarness.open();

        final ConcurrentLinkedQueue<Object> results = testHarness.getOutput();
        final AtomicBoolean throwOnCollect = new AtomicBoolean();

        @SuppressWarnings("unchecked")
        SourceFunction.SourceContext<String> sourceContext =
                new CollectingSourceContext(testHarness.getCheckpointLock(), results) {
                    @Override
                    public void markAsTemporarilyIdle() {}

                    @Override
                    public void collect(Serializable element) {
                        if (throwOnCollect.get()) {
                            throw new RuntimeException("expected");
                        }
                        super.collect(element);
                    }

                    @Override
                    public void emitWatermark(Watermark mark) {
                        results.add(mark);
                    }
                };

        final AtomicReference<Exception> sourceThreadError = new AtomicReference<>();
        new Thread(
                        () -> {
                            try {
                                sourceFunc.run(sourceContext);
                            } catch (InterruptedException e) {
                                // expected on cancel
                            } catch (Exception e) {
                                sourceThreadError.set(e);
                            }
                        })
                .start();

        ArrayList<Object> expectedResults = new ArrayList<>();

        final long record1 = 1;
        shard1.put(Long.toString(record1));
        expectedResults.add(Long.toString(record1));
        awaitRecordCount(results, expectedResults.size());

        // at this point we know the fetcher was initialized
        final KinesisDataFetcher fetcher =
                org.powermock.reflect.Whitebox.getInternalState(sourceFunc, "fetcher");

        // trigger watermark emit
        testHarness.setProcessingTime(testHarness.getProcessingTime() + autoWatermarkInterval);
        expectedResults.add(new Watermark(-4));
        // verify watermark
        awaitRecordCount(results, expectedResults.size());
        assertThat(results).contains(expectedResults.toArray());
        assertThat(TestWatermarkTracker.WATERMARK.get()).isEqualTo(0);

        // trigger sync
        testHarness.setProcessingTime(testHarness.getProcessingTime() + 1);
        TestWatermarkTracker.assertGlobalWatermark(-4);

        final long record2 = record1 + (watermarkSyncInterval * 3) + 1;
        shard1.put(Long.toString(record2));

        // wait for the record to be buffered in the emitter
        final RecordEmitter<?> emitter =
                org.powermock.reflect.Whitebox.getInternalState(fetcher, "recordEmitter");
        RecordEmitter.RecordQueue emitterQueue = emitter.getQueue(0);
        Deadline deadline = Deadline.fromNow(Duration.ofSeconds(10));
        while (deadline.hasTimeLeft() && emitterQueue.getSize() < 1) {
            Thread.sleep(10);
        }
        assertThat(emitterQueue.getSize()).as("first record received").isEqualTo(1);

        // Advance the watermark. Since the new record is past global watermark + threshold,
        // it won't be emitted and the watermark does not advance
        testHarness.setProcessingTime(testHarness.getProcessingTime() + autoWatermarkInterval);
        assertThat(results).contains(expectedResults.toArray());
        assertThat((long) org.powermock.reflect.Whitebox.getInternalState(fetcher, "nextWatermark"))
                .isEqualTo(3000L);
        TestWatermarkTracker.assertGlobalWatermark(-4);

        // Trigger global watermark sync
        testHarness.setProcessingTime(testHarness.getProcessingTime() + 1);
        expectedResults.add(Long.toString(record2));
        awaitRecordCount(results, expectedResults.size());
        assertThat(results).contains(expectedResults.toArray());
        TestWatermarkTracker.assertGlobalWatermark(3000);

        // Trigger watermark update and emit
        testHarness.setProcessingTime(testHarness.getProcessingTime() + autoWatermarkInterval);
        expectedResults.add(new Watermark(3000));
        assertThat(results).contains(expectedResults.toArray());

        // verify exception propagation
        assertThat(sourceThreadError.get()).isNull();
        throwOnCollect.set(true);
        shard1.put(Long.toString(record2 + 1));

        deadline = Deadline.fromNow(Duration.ofSeconds(10));
        while (deadline.hasTimeLeft() && sourceThreadError.get() == null) {
            Thread.sleep(10);
        }
        assertThat(sourceThreadError.get()).isNotNull();
        assertThat(sourceThreadError.get().getMessage()).as("expected").isNotNull();

        sourceFunc.cancel();
        testHarness.close();
    }

    @Test
    public void testCloseConnectorBeforeSubtaskStart() throws Exception {
        Properties config = TestUtils.getStandardProperties();
        FlinkKinesisConsumer<String> consumer =
                new FlinkKinesisConsumer<>("fakeStream", new SimpleStringSchema(), config);
        consumer.close();
    }

    private void awaitRecordCount(ConcurrentLinkedQueue<? extends Object> queue, int count)
            throws Exception {
        Deadline deadline = Deadline.fromNow(Duration.ofSeconds(10));
        while (deadline.hasTimeLeft() && queue.size() < count) {
            Thread.sleep(10);
        }

        int received = queue.size();
        assertThat(received).isEqualTo(count);
    }

    private static class OpenCheckingStringSchema extends SimpleStringSchema {
        private boolean opened = false;

        @Override
        public void open(DeserializationSchema.InitializationContext context) throws Exception {
            assertThat(context.getMetricGroup()).isNotNull();
            this.opened = true;
        }

        @Override
        public String deserialize(byte[] message) {
            if (!opened) {
                throw new AssertionError(
                        "DeserializationSchema was not opened before deserialization.");
            }
            return super.deserialize(message);
        }
    }

    private static class TestTimestampExtractor
            extends BoundedOutOfOrdernessTimestampExtractor<String> {
        private static final long serialVersionUID = 1L;

        public TestTimestampExtractor(Time maxAllowedLateness) {
            super(maxAllowedLateness);
        }

        @Override
        public long extractTimestamp(String element) {
            return Long.parseLong(element);
        }
    }

    private static class TestWatermarkTracker extends WatermarkTracker {

        private static final AtomicLong WATERMARK = new AtomicLong();

        @Override
        public long getUpdateTimeoutCount() {
            return 0;
        }

        @Override
        public long updateWatermark(long localWatermark) {
            WATERMARK.set(localWatermark);
            return localWatermark;
        }

        static void assertGlobalWatermark(long expected) {
            assertThat(WATERMARK.get()).isEqualTo(expected);
        }
    }
}