[hotfix][tests] Transition to RECOVERING before RUNNING in DefaultSch…

…edulerTest
apache · Apr 6, 2021 · 9cb1610 · 9cb1610
1 parent aadd02e
commit 9cb1610
Showing 1 changed file with 12 additions and 6 deletions.
diff --git a/flink-runtime/src/test/java/org/apache/flink/runtime/scheduler/DefaultSchedulerTest.java b/flink-runtime/src/test/java/org/apache/flink/runtime/scheduler/DefaultSchedulerTest.java
@@ -744,8 +744,7 @@ public void abortPendingCheckpointsWhenRestartingTasks() throws Exception {
                                 .getAllExecutionVertices());
         final ExecutionAttemptID attemptId =
                 onlyExecutionVertex.getCurrentExecutionAttempt().getAttemptId();
-        scheduler.updateTaskExecutionState(
-                new TaskExecutionState(attemptId, ExecutionState.RUNNING));
+        transitionToRunning(scheduler, attemptId);
 
         final CheckpointCoordinator checkpointCoordinator = getCheckpointCoordinator(scheduler);
 
@@ -775,8 +774,7 @@ public void restoreStateWhenRestartingTasks() throws Exception {
                                 .getAllExecutionVertices());
         final ExecutionAttemptID attemptId =
                 onlyExecutionVertex.getCurrentExecutionAttempt().getAttemptId();
-        scheduler.updateTaskExecutionState(
-                new TaskExecutionState(attemptId, ExecutionState.RUNNING));
+        transitionToRunning(scheduler, attemptId);
 
         final CheckpointCoordinator checkpointCoordinator = getCheckpointCoordinator(scheduler);
 
@@ -814,8 +812,7 @@ public void failGlobalWhenRestoringStateFails() throws Exception {
                                 .getAllExecutionVertices());
         final ExecutionAttemptID attemptId =
                 onlyExecutionVertex.getCurrentExecutionAttempt().getAttemptId();
-        scheduler.updateTaskExecutionState(
-                new TaskExecutionState(attemptId, ExecutionState.RUNNING));
+        transitionToRunning(scheduler, attemptId);
 
         final CheckpointCoordinator checkpointCoordinator = getCheckpointCoordinator(scheduler);
 
@@ -1326,4 +1323,13 @@ private CountDownLatch getCheckpointTriggeredLatch() {
                 });
         return checkpointTriggeredLatch;
     }
+
+    private void transitionToRunning(DefaultScheduler scheduler, ExecutionAttemptID attemptId) {
+        Preconditions.checkState(
+                scheduler.updateTaskExecutionState(
+                        new TaskExecutionState(attemptId, ExecutionState.RECOVERING)));
+        Preconditions.checkState(
+                scheduler.updateTaskExecutionState(
+                        new TaskExecutionState(attemptId, ExecutionState.RUNNING)));
+    }
 }