/*
 * 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.base.source.reader.fetcher;

import org.apache.flink.api.connector.source.mocks.MockSourceSplit;
import org.apache.flink.connector.base.source.reader.RecordsWithSplitIds;
import org.apache.flink.connector.base.source.reader.mocks.MockSplitReader;
import org.apache.flink.connector.base.source.reader.synchronization.FutureCompletingBlockingQueue;
import org.apache.flink.connector.base.source.reader.synchronization.FutureNotifier;

import org.junit.Test;

import java.util.ArrayList;
import java.util.Collections;
import java.util.List;
import java.util.SortedSet;
import java.util.TreeSet;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.concurrent.atomic.AtomicInteger;

import static org.junit.Assert.assertEquals;
import static org.junit.Assert.assertTrue;

/**
 * Unit test for {@link SplitFetcher}.
 */
public class SplitFetcherTest {
	private static final int NUM_SPLITS = 3;
	private static final int NUM_RECORDS_PER_SPLIT = 10_000;
	private static final int INTERRUPT_RECORDS_INTERVAL = 10;
	private static final int NUM_TOTAL_RECORDS = NUM_RECORDS_PER_SPLIT * NUM_SPLITS;

	@Test
	public void testWakeup() throws InterruptedException {
		FutureCompletingBlockingQueue<RecordsWithSplitIds<int[]>> elementQueue =
			new FutureCompletingBlockingQueue<>(new FutureNotifier(), 1);
		SplitFetcher<int[], MockSourceSplit> fetcher =
				new SplitFetcher<>(
						0,
						elementQueue,
						new MockSplitReader(2, true, true),
						() -> {});

		// Prepare the splits.
		List<MockSourceSplit> splits = new ArrayList<>();
		for (int i = 0; i < NUM_SPLITS; i++) {
			splits.add(new MockSourceSplit(i, 0, NUM_RECORDS_PER_SPLIT));
			int base = i * NUM_RECORDS_PER_SPLIT;
			for (int j = base; j < base + NUM_RECORDS_PER_SPLIT; j++) {
				splits.get(splits.size() - 1).addRecord(j);
			}
		}
		// Add splits to the fetcher.
		fetcher.addSplits(splits);

		// A thread drives the fetcher.
		Thread fetcherThread = new Thread(fetcher, "FetcherThread");

		SortedSet<Integer> recordsRead = Collections.synchronizedSortedSet(new TreeSet<>());

		// A thread waking up the split fetcher frequently.
		AtomicInteger wakeupTimes = new AtomicInteger(0);
		AtomicBoolean stop = new AtomicBoolean(false);
		Thread interrupter = new Thread("Interrupter") {
			@Override
			public void run() {
				int lastInterrupt = 0;
				while (recordsRead.size() < NUM_TOTAL_RECORDS && !stop.get()) {
					int numRecordsRead = recordsRead.size();
					if (numRecordsRead >= lastInterrupt + INTERRUPT_RECORDS_INTERVAL) {
						fetcher.wakeUp(false);
						wakeupTimes.incrementAndGet();
						lastInterrupt = numRecordsRead;
					}
				}
			}
		};

		try {
			fetcherThread.start();
			interrupter.start();

			while (recordsRead.size() < NUM_SPLITS * NUM_RECORDS_PER_SPLIT) {
				final RecordsWithSplitIds<int[]> nextBatch = elementQueue.take();
				while (nextBatch.nextSplit() != null) {
					int[] arr;
					while ((arr = nextBatch.nextRecordFromSplit()) != null) {
						assertTrue(recordsRead.add(arr[0]));
					}
				}
			}

			assertEquals(NUM_TOTAL_RECORDS, recordsRead.size());
			assertEquals(0, (int) recordsRead.first());
			assertEquals(NUM_TOTAL_RECORDS - 1, (int) recordsRead.last());
			assertTrue(wakeupTimes.get() > 0);
		} finally {
			stop.set(true);
			fetcher.shutdown();
			fetcherThread.join();
			interrupter.join();
		}
	}
}

@@ -1700,312 +1700,245 @@
 t %7B%0A
-%09priva
+%0A%09@Test%0A%09public void 
 te
- 
 st
-atic final int NUM_SPLITS = 3;%0A%09private static final int NUM_RECORDS_PER_SPLIT = 10_000;%0A%09private static final int INTERRUPT_RECORDS_INTERVAL = 10;%0A%09private static final int NUM_TOTAL_RECORDS = NUM_RECORDS_PER_SPLIT * NUM_SPLITS;%0A%0A%09@Test%0A%09public void testWakeup() throws InterruptedException %7B
+Wakeup() throws InterruptedException %7B%0A%09%09final int numSplits = 3;%0A%09%09final int numRecordsPerSplit = 10_000;%0A%09%09final int interruptRecordsInterval = 10;%0A%09%09final int numTotalRecords = numRecordsPerSplit * numSplits;%0A
 %0A%09%09F
@@ -2335,26 +2335,25 @@
  0; i %3C 
-NUM_SPLITS
+numSplits
 ; i++) %7B
@@ -2393,37 +2393,34 @@
 t(i, 0, 
-NUM_RECORDS_PER_SPLIT
+numRecordsPerSplit
 ));%0A%09%09%09i
@@ -2433,37 +2433,34 @@
 e = i * 
-NUM_RECORDS_PER_SPLIT
+numRecordsPerSplit
 ;%0A%09%09%09for
@@ -2486,37 +2486,34 @@
  base + 
-NUM_RECORDS_PER_SPLIT
+numRecordsPerSplit
 ; j++) %7B
@@ -3119,33 +3119,31 @@
 ize() %3C 
-NUM_TOTAL_RECORDS
+numTotalRecords
  && !sto
@@ -3246,34 +3246,32 @@
 t + 
-INTERRUPT_RECORDS_INTERVAL
+interruptRecordsInterval
 ) %7B%0A
@@ -3492,42 +3492,38 @@
 ) %3C 
-NUM_SPLITS * NUM_RECORDS_PER_SPLIT
+numSplits * numRecordsPerSplit
 ) %7B%0A
@@ -3793,33 +3793,31 @@
 tEquals(
-NUM_TOTAL_RECORDS
+numTotalRecords
 , record
@@ -3898,25 +3898,23 @@
 als(
-NUM_TOTAL_RECORDS
+numTotalRecords
  - 1

Move constants in SplitFetcherTest relevant to only one test into test method