4343import org .apache .flink .api .connector .source .SourceEvent ;
4444import org .apache .flink .api .connector .source .SplitsAssignment ;
4545import org .apache .flink .api .connector .source .mocks .MockSplitEnumeratorContext ;
46+ import org .jetbrains .annotations .NotNull ;
4647import org .junit .jupiter .api .BeforeAll ;
4748import org .junit .jupiter .api .Test ;
4849import org .junit .jupiter .params .ParameterizedTest ;
5859import java .util .List ;
5960import java .util .Map ;
6061import java .util .Set ;
62+ import java .util .concurrent .Callable ;
63+ import java .util .concurrent .ScheduledExecutorService ;
64+ import java .util .concurrent .ScheduledFuture ;
65+ import java .util .concurrent .TimeUnit ;
6166
6267import static org .apache .fluss .client .table .scanner .log .LogScanner .EARLIEST_OFFSET ;
6368import static org .apache .fluss .testutils .DataTestUtils .row ;
@@ -97,7 +102,8 @@ void testPkTableNoSnapshotSplits() throws Throwable {
97102 OffsetsInitializer .full (),
98103 DEFAULT_SCAN_PARTITION_DISCOVERY_INTERVAL_MS ,
99104 streaming ,
100- Collections .emptyList ());
105+ Collections .emptyList (),
106+ null );
101107
102108 enumerator .start ();
103109
@@ -144,7 +150,7 @@ void testPkTableWithSnapshotSplits() throws Throwable {
144150 OffsetsInitializer .full (),
145151 DEFAULT_SCAN_PARTITION_DISCOVERY_INTERVAL_MS ,
146152 streaming ,
147- Collections .emptyList ());
153+ Collections .emptyList (), null );
148154 enumerator .start ();
149155 // register all read
150156 for (int i = 0 ; i < numSubtasks ; i ++) {
@@ -215,7 +221,7 @@ void testNonPkTable() throws Throwable {
215221 OffsetsInitializer .full (),
216222 DEFAULT_SCAN_PARTITION_DISCOVERY_INTERVAL_MS ,
217223 streaming ,
218- Collections .emptyList ());
224+ Collections .emptyList (), null );
219225
220226 enumerator .start ();
221227
@@ -261,7 +267,7 @@ void testReaderRegistrationTriggerAssignments() throws Throwable {
261267 OffsetsInitializer .full (),
262268 DEFAULT_SCAN_PARTITION_DISCOVERY_INTERVAL_MS ,
263269 streaming ,
264- Collections .emptyList ());
270+ Collections .emptyList (), null );
265271
266272 enumerator .start ();
267273
@@ -297,7 +303,7 @@ void testAddSplitBack() throws Throwable {
297303 OffsetsInitializer .full (),
298304 DEFAULT_SCAN_PARTITION_DISCOVERY_INTERVAL_MS ,
299305 streaming ,
300- Collections .emptyList ());
306+ Collections .emptyList (), null );
301307
302308 enumerator .start ();
303309
@@ -339,6 +345,7 @@ void testRestore() throws Throwable {
339345 try (MockSplitEnumeratorContext <SourceSplitBase > context =
340346 new MockSplitEnumeratorContext <>(numSubtasks )) {
341347
348+
342349 // mock bucket1 has been assigned
343350 TableBucket bucket1 = new TableBucket (tableId , 1 );
344351 Set <TableBucket > assignedBuckets = new HashSet <>(Collections .singletonList (bucket1 ));
@@ -391,30 +398,35 @@ void testDiscoverPartitionsPeriodically(boolean isPrimaryKeyTable) throws Throwa
391398 ZooKeeperClient zooKeeperClient = FLUSS_CLUSTER_EXTENSION .getZooKeeperClient ();
392399 try (MockSplitEnumeratorContext <SourceSplitBase > context =
393400 new MockSplitEnumeratorContext <>(numSubtasks );
401+ MockWorkExecutor workExecutor = new MockWorkExecutor (context );
394402 FlinkSourceEnumerator enumerator =
395403 new FlinkSourceEnumerator (
396404 DEFAULT_TABLE_PATH ,
397405 flussConf ,
398406 isPrimaryKeyTable ,
399407 true ,
400408 context ,
409+ Collections .emptySet (),
410+ Collections .emptyMap (),
411+ null ,
401412 OffsetsInitializer .full (),
402413 DEFAULT_SCAN_PARTITION_DISCOVERY_INTERVAL_MS ,
403414 streaming ,
404- Collections .emptyList ())) {
415+ Collections .emptyList (), null ,
416+ workExecutor )) {
405417 Map <Long , String > partitionNameByIds =
406418 waitUntilPartitions (zooKeeperClient , DEFAULT_TABLE_PATH );
407419 enumerator .start ();
408420
409421 // invoke partition discovery callable again and there should be pending assignments.
410- runPeriodicPartitionDiscovery (context );
422+ runPeriodicPartitionDiscovery ( workExecutor );
411423
412424 // register two readers
413425 registerReader (context , enumerator , 0 );
414426 registerReader (context , enumerator , 1 );
415427
416428 // invoke partition discovery callable again, shouldn't produce RemovePartitionEvent.
417- runPeriodicPartitionDiscovery (context );
429+ runPeriodicPartitionDiscovery ( workExecutor );
418430 assertThat (context .getSentSourceEvent ()).isEmpty ();
419431
420432 // now, register the third reader
@@ -434,7 +446,7 @@ void testDiscoverPartitionsPeriodically(boolean isPrimaryKeyTable) throws Throwa
434446 createPartitions (zooKeeperClient , DEFAULT_TABLE_PATH , newPartitions );
435447
436448 /// invoke partition discovery callable again and there should assignments.
437- runPeriodicPartitionDiscovery (context );
449+ runPeriodicPartitionDiscovery (workExecutor );
438450
439451 expectedAssignment = expectAssignments (enumerator , tableId , newPartitionNameIds );
440452 actualAssignments = getLastReadersAssignments (context );
@@ -450,7 +462,7 @@ void testDiscoverPartitionsPeriodically(boolean isPrimaryKeyTable) throws Throwa
450462 createPartitions (zooKeeperClient , DEFAULT_TABLE_PATH , newPartitions );
451463
452464 // invoke partition discovery callable again
453- runPeriodicPartitionDiscovery (context );
465+ runPeriodicPartitionDiscovery ( workExecutor );
454466
455467 // there should be partition removed events
456468 Map <Integer , List <SourceEvent >> sentSourceEvents = context .getSentSourceEvent ();
@@ -516,7 +528,7 @@ void testGetSplitOwner() throws Exception {
516528 OffsetsInitializer .full (),
517529 DEFAULT_SCAN_PARTITION_DISCOVERY_INTERVAL_MS ,
518530 streaming ,
519- Collections .emptyList ())) {
531+ Collections .emptyList (), null )) {
520532
521533 // test splits for same non-partitioned bucket, should assign to same task
522534 TableBucket t1 = new TableBucket (tableId , 0 );
@@ -587,13 +599,13 @@ private void checkAssignmentIgnoreOrder(
587599 }
588600 }
589601
590- private void runPeriodicPartitionDiscovery (MockSplitEnumeratorContext < SourceSplitBase > context )
602+ private void runPeriodicPartitionDiscovery ( MockWorkExecutor workExecutor )
591603 throws Throwable {
592604 // Fetch potential topic descriptions
593- context .runPeriodicCallable (PARTITION_DISCOVERY_CALLABLE_INDEX );
605+ workExecutor .runPeriodicCallable (PARTITION_DISCOVERY_CALLABLE_INDEX );
594606 // Initialize offsets for discovered partitions
595- if (!context .getOneTimeCallables ().isEmpty ()) {
596- context .runNextOneTimeCallable ();
607+ if (!workExecutor .getOneTimeCallables ().isEmpty ()) {
608+ workExecutor .runNextOneTimeCallable ();
597609 }
598610 }
599611
0 commit comments