Fix test failure

Author: yunfengzhou-hub

paimon-flink/paimon-flink-common/src/test/java/org/apache/paimon/flink/FlinkChainTableITCase.java

@@ -31,12 +31,9 @@
 import org.apache.paimon.table.source.TableScan;
 import org.apache.paimon.utils.BlockingIterator;
 
-import org.apache.flink.api.common.JobID;
 import org.apache.flink.configuration.CheckpointingOptions;
 import org.apache.flink.configuration.ExternalizedCheckpointRetention;
 import org.apache.flink.core.execution.JobClient;
-import org.apache.flink.runtime.execution.ExecutionState;
-import org.apache.flink.runtime.minicluster.MiniCluster;
 import org.apache.flink.streaming.api.datastream.DataStream;
 import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
 import org.apache.flink.table.api.DataTypes;
@@ -47,12 +44,10 @@
 import org.junit.jupiter.api.Test;
 import org.junit.jupiter.api.Timeout;
 
-import java.lang.reflect.Field;
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.List;
 import java.util.concurrent.TimeUnit;
-import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.stream.Collectors;
 
 import static org.assertj.core.api.Assertions.assertThat;
@@ -781,37 +776,6 @@ private void waitForRowCount(String tableName, int minRows) throws Exception {
                 "Timed out waiting for " + minRows + " rows in " + tableName + ", got " + count);
     }
 
-    /** Polls until all tasks of the given job are in RUNNING state. */
-    private void waitForJobRunning(JobClient jobClient) throws Exception {
-        Field field = jobClient.getClass().getDeclaredField("miniCluster");
-        field.setAccessible(true);
-        MiniCluster miniCluster = (MiniCluster) field.get(jobClient);
-        JobID jobID = jobClient.getJobID();
-
-        long deadline = System.currentTimeMillis() + 60_000;
-        while (System.currentTimeMillis() < deadline) {
-            AtomicBoolean allRunning = new AtomicBoolean(true);
-            miniCluster
-                    .getExecutionGraph(jobID)
-                    .thenAccept(
-                            eg ->
-                                    eg.getAllExecutionVertices()
-                                            .forEach(
-                                                    v -> {
-                                                        if (v.getExecutionState()
-                                                                != ExecutionState.RUNNING) {
-                                                            allRunning.set(false);
-                                                        }
-                                                    }))
-                    .get();
-            if (allRunning.get()) {
-                return;
-            }
-            Thread.sleep(1000);
-        }
-        throw new AssertionError("Timed out waiting for job " + jobID + " to reach RUNNING state");
-    }
-
     /**
      * Tests the streaming read lifecycle for a chain table with changelog-producer=input.
      *
@@ -1051,9 +1015,6 @@ public void testStreamingReadChainTableStatefulRestart() throws Exception {
         //noinspection OptionalGetWithoutIsPresent
         JobClient jobClient = tableResult.getJobClient().get();
 
-        // Wait for streaming job to be fully running before writing Phase 2 data.
-        waitForJobRunning(jobClient);
-
         // === Phase 2: Write incremental delta, let Phase 2 consume it, THEN checkpoint ===
         // This exercises the checkpoint() regression: if checkpoint() returns the stale
         // Phase 1 boundary instead of the advanced delta cursor, restore would re-read
@@ -1067,16 +1028,18 @@ public void testStreamingReadChainTableStatefulRestart() throws Exception {
         // advanced cursor — the exact scenario the regression would break.
         waitForRowCount("chain_restart_sink", 7);
 
-        // Create a savepoint for restart. The enumerator state now includes the advanced
-        // delta cursor (past delta@20250809).
-        String checkpointPath = triggerCheckpoint(jobClient);
-
-        java.io.File cpFile =
-                new java.io.File(checkpointPath.replace("file:/", "/").replace("file://", "/"));
-        if (!cpFile.exists()) {
-            cpFile = cpFile.getParentFile();
-        }
-        assertThat(cpFile.exists()).as("Checkpoint directory should exist: " + cpFile).isTrue();
+        // Create a savepoint and stop the job atomically.
+        // stopWithSavepoint guarantees the savepoint is fully committed to disk
+        // before returning.
+        String savepointDir = path + "/savepoints";
+        new java.io.File(savepointDir).mkdirs();
+        String checkpointPath =
+                jobClient
+                        .stopWithSavepoint(
+                                false,
+                                savepointDir,
+                                org.apache.flink.core.execution.SavepointFormatType.CANONICAL)
+                        .get();
 
         // Verify Phase 1+2 data (committed by checkpoint).
         List<String> phase1and2 =
@@ -1094,20 +1057,14 @@ public void testStreamingReadChainTableStatefulRestart() throws Exception {
                         "+I[4, 1, new_4, 20250809]",
                         "+I[5, 1, new_5, 20250809]");
 
-        jobClient.cancel().get();
-
-        // Disable auto-checkpointing before restart to avoid conflicts between
-        // auto-checkpoints and manual triggerCheckpoint() on the new job.
-        config.removeKey("execution.checkpointing.interval");
-
         // === Phase 3: Write new delta data while job is stopped ===
         sql(
                 "INSERT INTO `chain_restart$branch_delta` PARTITION (dt = '20250810')"
                         + " VALUES (6, 1, 'new_6'), (7, 1, 'new_7')");
 
-        // === Phase 4: Restart from checkpoint ===
+        // === Phase 4: Restart from savepoint ===
         // The restored scan should NOT re-read delta@20250809 (already consumed before
-        // checkpoint). If checkpoint() returned the stale Phase 1 boundary, delta@20250809
+        // savepoint). If checkpoint() returned the stale Phase 1 boundary, delta@20250809
         // would be re-read and produce duplicates.
         sEnv.getConfig()
                 .getConfiguration()
@@ -1117,25 +1074,19 @@ public void testStreamingReadChainTableStatefulRestart() throws Exception {
         //noinspection OptionalGetWithoutIsPresent
         JobClient jobClient2 = tableResult2.getJobClient().get();
 
-        // Trigger checkpoint to commit Phase 4 data (restored scan output + Phase 3 delta).
-        waitForJobRunning(jobClient2);
-        triggerCheckpoint(jobClient2);
-
-        // Poll sink until checkpoint commits all 9 rows (7 from Phase 1+2 + 2 from Phase 3).
+        // Auto-checkpointing (1s interval) is still enabled, so data is committed
+        // to the sink automatically. waitForRowCount polls until data appears.
         waitForRowCount("chain_restart_sink", 9);
 
-        // Read all records from the sink table
         List<String> phase4 =
                 sql("SELECT * FROM chain_restart_sink").stream()
                         .map(Row::toString)
                         .collect(Collectors.toList());
 
-        // Verify new data is present
         assertThat(phase4)
                 .as("Stateful restart: sink should contain new delta data")
                 .contains("+I[6, 1, new_6, 20250810]", "+I[7, 1, new_7, 20250810]");
 
-        // Verify total count: 7 from Phase 1+2 + 2 from Phase 3 = 9, no duplicates
         assertThat(phase4.size())
                 .as("Should have exactly 9 records (no duplicates from state recovery)")
                 .isEqualTo(9);
@@ -1145,9 +1096,6 @@ public void testStreamingReadChainTableStatefulRestart() throws Exception {
                 "INSERT INTO `chain_restart$branch_delta` PARTITION (dt = '20250811')"
                         + " VALUES (8, 1, 'new_8')");
 
-        triggerCheckpoint(jobClient2);
-
-        // Poll sink until checkpoint commits the new incremental row.
         waitForRowCount("chain_restart_sink", 10);
 
         List<String> phase5 =
@@ -1461,39 +1409,6 @@ public void testStreamingReadWithGroupPartition() throws Exception {
         it.close();
     }
 
-    /**
-     * Triggers a checkpoint on the given job and returns the checkpoint path. Waits until all tasks
-     * are RUNNING before triggering.
-     */
-    private String triggerCheckpoint(JobClient jobClient) throws Exception {
-        Field field = jobClient.getClass().getDeclaredField("miniCluster");
-        field.setAccessible(true);
-        MiniCluster miniCluster = (MiniCluster) field.get(jobClient);
-        JobID jobID = jobClient.getJobID();
-
-        // Wait for all tasks to be RUNNING
-        AtomicBoolean allRunning = new AtomicBoolean(false);
-        while (!allRunning.get()) {
-            allRunning.set(true);
-            Thread.sleep(1000);
-            miniCluster
-                    .getExecutionGraph(jobID)
-                    .thenAccept(
-                            eg ->
-                                    eg.getAllExecutionVertices()
-                                            .forEach(
-                                                    v -> {
-                                                        if (v.getExecutionState()
-                                                                != ExecutionState.RUNNING) {
-                                                            allRunning.set(false);
-                                                        }
-                                                    }))
-                    .get();
-        }
-
-        return miniCluster.triggerCheckpoint(jobID).get();
-    }
-
     // =========================================================================
     // Additional coverage tests
     // =========================================================================