[flink] Cover restore through FlinkSource

Author: luoyuxia

fluss-flink/fluss-flink-common/src/test/java/org/apache/fluss/flink/source/enumerator/FlinkSourceEnumeratorTest.java

@@ -248,48 +248,52 @@ void testRestoreFlussOnlySourceWithLakeSourceDoesNotGenerateLakeSplits(@TempDir
 
         SourceEnumeratorState checkpointState;
         try (MockSplitEnumeratorContext<SourceSplitBase> context =
-                new MockSplitEnumeratorContext<>(1)) {
-            FlinkSourceEnumerator enumerator =
-                    new FlinkSourceEnumerator(
-                            DEFAULT_TABLE_PATH,
-                            flussConf,
-                            true,
-                            false,
-                            context,
-                            OffsetsInitializer.timestamp(1000L),
-                            DEFAULT_SCAN_PARTITION_DISCOVERY_INTERVAL_MS,
-                            streaming,
-                            null,
-                            null,
-                            LeaseContext.DEFAULT,
-                            false);
-
+                        new MockSplitEnumeratorContext<>(1);
+                SplitEnumerator<SourceSplitBase, SourceEnumeratorState> enumerator =
+                        new FlinkSource<RowData>(
+                                        flussConf,
+                                        DEFAULT_TABLE_PATH,
+                                        false,
+                                        true,
+                                        DEFAULT_LOG_TABLE_SCHEMA.getRowType(),
+                                        null,
+                                        null,
+                                        OffsetsInitializer.timestamp(1000L),
+                                        0L,
+                                        new RowDataDeserializationSchema(),
+                                        streaming,
+                                        null,
+                                        LeaseContext.DEFAULT)
+                                .createEnumerator(context)) {
             checkpointState = enumerator.snapshotState(1L);
+            assertThat(checkpointState.getRemainingHybridLakeFlussSplits()).isNull();
         }
 
         try (MockSplitEnumeratorContext<SourceSplitBase> context =
                         new MockSplitEnumeratorContext<>(DEFAULT_BUCKET_NUM);
-                MockWorkExecutor workExecutor = new MockWorkExecutor(context);
-                FlinkSourceEnumerator restoredEnumerator =
-                        new FlinkSourceEnumerator(
-                                DEFAULT_TABLE_PATH,
-                                flussConf,
-                                false,
-                                true,
-                                context,
-                                checkpointState.getAssignedBuckets(),
-                                checkpointState.getAssignedPartitions(),
-                                Collections.emptyList(),
-                                OffsetsInitializer.full(),
-                                DEFAULT_SCAN_PARTITION_DISCOVERY_INTERVAL_MS,
-                                streaming,
-                                null,
-                                lakeSource,
-                                workExecutor,
-                                LeaseContext.DEFAULT,
-                                true)) {
+                SplitEnumerator<SourceSplitBase, SourceEnumeratorState> restoredEnumerator =
+                        new FlinkSource<RowData>(
+                                        flussConf,
+                                        DEFAULT_TABLE_PATH,
+                                        false,
+                                        true,
+                                        DEFAULT_LOG_TABLE_SCHEMA.getRowType(),
+                                        null,
+                                        null,
+                                        OffsetsInitializer.full(),
+                                        0L,
+                                        new RowDataDeserializationSchema(),
+                                        streaming,
+                                        null,
+                                        lakeSource,
+                                        LeaseContext.DEFAULT)
+                                .restoreEnumerator(context, checkpointState)) {
+            assertThat(restoredEnumerator.snapshotState(1L).getRemainingHybridLakeFlussSplits())
+                    .isEmpty();
+
             restoredEnumerator.start();
-            runPeriodicPartitionDiscovery(workExecutor);
+            context.runNextOneTimeCallable();
+            context.runNextOneTimeCallable();
 
             for (int i = 0; i < DEFAULT_BUCKET_NUM; i++) {
                 registerReader(context, restoredEnumerator, i);
@@ -305,40 +309,6 @@ void testRestoreFlussOnlySourceWithLakeSourceDoesNotGenerateLakeSplits(@TempDir
         }
     }
 
-    @Test
-    void testRestoreEnumeratorNormalizesNullLakeSplits() throws Exception {
-        SourceEnumeratorState checkpointState =
-                new SourceEnumeratorState(
-                        Collections.emptySet(),
-                        Collections.emptyMap(),
-                        null,
-                        LeaseContext.DEFAULT.getKvSnapshotLeaseId());
-        FlinkSource<RowData> source =
-                new FlinkSource<>(
-                        flussConf,
-                        DEFAULT_TABLE_PATH,
-                        false,
-                        true,
-                        DEFAULT_LOG_TABLE_SCHEMA.getRowType(),
-                        null,
-                        null,
-                        OffsetsInitializer.full(),
-                        DEFAULT_SCAN_PARTITION_DISCOVERY_INTERVAL_MS,
-                        new RowDataDeserializationSchema(),
-                        streaming,
-                        null,
-                        new TestingLakeSource(DEFAULT_BUCKET_NUM, Collections.emptyList()),
-                        LeaseContext.DEFAULT);
-
-        try (MockSplitEnumeratorContext<SourceSplitBase> context =
-                        new MockSplitEnumeratorContext<>(1);
-                SplitEnumerator<SourceSplitBase, SourceEnumeratorState> restoredEnumerator =
-                        source.restoreEnumerator(context, checkpointState)) {
-            assertThat(restoredEnumerator.snapshotState(1L).getRemainingHybridLakeFlussSplits())
-                    .isEmpty();
-        }
-    }
-
     @Test
     void testPkTableWithSnapshotSplits() throws Throwable {
         long tableId = createTable(DEFAULT_TABLE_PATH, DEFAULT_PK_TABLE_DESCRIPTOR);
@@ -981,7 +951,7 @@ void testPartitionsExpiredInFlussButExistInLake(
     // ---------------------
     private void registerReader(
             MockSplitEnumeratorContext<SourceSplitBase> context,
-            FlinkSourceEnumerator enumerator,
+            SplitEnumerator<SourceSplitBase, SourceEnumeratorState> enumerator,
             int readerId) {
         context.registerReader(new ReaderInfo(readerId, "location " + readerId));
         enumerator.addReader(readerId);