[FLINK-32695] Migrated SavepointReaderITTestBase to SourceV2

Author: Poorvankbhatia

flink-libraries/flink-state-processing-api/pom.xml

@@ -114,6 +114,13 @@ under the License.
 			<scope>test</scope>
 		</dependency>
 
+		<dependency>
+			<groupId>org.apache.flink</groupId>
+			<artifactId>flink-test-utils-connector</artifactId>
+			<version>${project.version}</version>
+			<scope>test</scope>
+		</dependency>
+
 	</dependencies>
 
 	<build>

flink-libraries/flink-state-processing-api/src/test/java/org/apache/flink/state/api/SavepointReaderITTestBase.java

@@ -19,14 +19,20 @@
 package org.apache.flink.state.api;
 
 import org.apache.flink.api.common.JobID;
+import org.apache.flink.api.common.eventtime.WatermarkStrategy;
 import org.apache.flink.api.common.functions.OpenContext;
 import org.apache.flink.api.common.state.ListState;
 import org.apache.flink.api.common.state.ListStateDescriptor;
 import org.apache.flink.api.common.state.MapStateDescriptor;
 import org.apache.flink.api.common.time.Deadline;
+import org.apache.flink.api.connector.source.ReaderOutput;
+import org.apache.flink.api.connector.source.SourceReaderContext;
+import org.apache.flink.api.connector.source.SplitEnumerator;
+import org.apache.flink.api.connector.source.SplitEnumeratorContext;
 import org.apache.flink.api.java.tuple.Tuple2;
 import org.apache.flink.client.program.ClusterClient;
 import org.apache.flink.core.execution.SavepointFormatType;
+import org.apache.flink.core.io.InputStatus;
 import org.apache.flink.runtime.jobgraph.JobGraph;
 import org.apache.flink.runtime.state.FunctionInitializationContext;
 import org.apache.flink.runtime.state.FunctionSnapshotContext;
@@ -37,8 +43,11 @@
 import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
 import org.apache.flink.streaming.api.functions.co.BroadcastProcessFunction;
 import org.apache.flink.streaming.api.functions.sink.v2.DiscardingSink;
-import org.apache.flink.streaming.api.functions.source.legacy.SourceFunction;
 import org.apache.flink.test.util.AbstractTestBaseJUnit4;
+import org.apache.flink.test.util.source.AbstractTestSource;
+import org.apache.flink.test.util.source.SingleSplitEnumerator;
+import org.apache.flink.test.util.source.TestSourceReader;
+import org.apache.flink.test.util.source.TestSplit;
 import org.apache.flink.util.AbstractID;
 import org.apache.flink.util.Collector;
 
@@ -88,7 +97,12 @@ public void testOperatorStateInputFormat() throws Exception {
         StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();
         env.setParallelism(4);
 
-        DataStream<Integer> data = env.addSource(new SavepointSource()).rebalance();
+        DataStream<Integer> data =
+                env.fromSource(
+                                new SavepointSource(),
+                                WatermarkStrategy.noWatermarks(),
+                                "SavepointSource")
+                        .rebalance();
 
         StatefulOperator statefulOperator = new StatefulOperator(list, union, broadcast);
         data.connect(data.broadcast(broadcast))
@@ -209,35 +223,50 @@ private String takeSavepoint(JobGraph jobGraph) throws Exception {
         }
     }
 
-    private static class SavepointSource implements SourceFunction<Integer> {
+    private static class SavepointSource extends AbstractTestSource<Integer> {
         private static volatile boolean finished;
 
-        private volatile boolean running = true;
-
         private static final Integer[] elements = {1, 2, 3};
 
         @Override
-        public void run(SourceContext<Integer> ctx) {
-            synchronized (ctx.getCheckpointLock()) {
-                for (Integer element : elements) {
-                    ctx.collect(element);
-                }
+        public SplitEnumerator<TestSplit, Void> createEnumerator(
+                SplitEnumeratorContext<TestSplit> enumContext) {
+            return new SingleSplitEnumerator(enumContext);
+        }
 
-                finished = true;
-            }
+        @Override
+        public TestSourceReader<Integer> createReader(SourceReaderContext readerContext) {
+            return new TestSourceReader<>(readerContext) {
+                private boolean receivedSplit = false;
+                private CompletableFuture<Void> availability = new CompletableFuture<>();
+
+                @Override
+                public void addSplits(List<TestSplit> splits) {
+                    if (!splits.isEmpty()) {
+                        receivedSplit = true;
+                        if (!availability.isDone()) {
+                            availability.complete(null);
+                        }
+                    }
+                }
 
-            while (running) {
-                try {
-                    Thread.sleep(100);
-                } catch (InterruptedException e) {
-                    // ignore
+                @Override
+                public InputStatus pollNext(ReaderOutput<Integer> output) {
+                    if (receivedSplit) {
+                        for (Integer element : elements) {
+                            output.collect(element);
+                        }
+                        finished = true;
+                        availability = new CompletableFuture<>();
+                    }
+                    return InputStatus.NOTHING_AVAILABLE;
                 }
-            }
-        }
 
-        @Override
-        public void cancel() {
-            running = false;
+                @Override
+                public CompletableFuture<Void> isAvailable() {
+                    return availability;
+                }
+            };
         }
 
         private static void initializeForTest() {