[FLINK-11343][tests] Force order that TaskExecutor shutdown after Tas…

…k exit in TaskExecutorTest

flink-runtime/src/test/java/org/apache/flink/runtime/taskexecutor/TaskExecutorTest.java

@@ -88,7 +88,9 @@
 import org.apache.flink.runtime.taskexecutor.slot.TimerService;
 import org.apache.flink.runtime.taskmanager.CheckpointResponder;
 import org.apache.flink.runtime.taskmanager.LocalTaskManagerLocation;
+import org.apache.flink.runtime.taskmanager.NoOpTaskManagerActions;
 import org.apache.flink.runtime.taskmanager.Task;
+import org.apache.flink.runtime.taskmanager.TaskExecutionState;
 import org.apache.flink.runtime.taskmanager.TaskManagerActions;
 import org.apache.flink.runtime.taskmanager.TaskManagerLocation;
 import org.apache.flink.runtime.testingUtils.TestingUtils;
@@ -715,11 +717,20 @@ public void testTaskSubmission() throws Exception {
  final JobMasterGateway jobMasterGateway = mock(JobMasterGateway.class);
  when(jobMasterGateway.getFencingToken()).thenReturn(jobMasterId);
 
+ final OneShotLatch taskInTerminalState = new OneShotLatch();
+ final TaskManagerActions taskManagerActions = new NoOpTaskManagerActions() {
+ @Override
+ public void updateTaskExecutionState(TaskExecutionState taskExecutionState) {
+ if (taskExecutionState.getExecutionState().isTerminal()) {
+ taskInTerminalState.trigger();
+ }
+ }
+ };
  final JobManagerConnection jobManagerConnection = new JobManagerConnection(
  jobId,
  ResourceID.generate(),
  jobMasterGateway,
- mock(TaskManagerActions.class),
+ taskManagerActions,
  mock(CheckpointResponder.class),
  new TestGlobalAggregateManager(),
  libraryCacheManager,
@@ -769,6 +780,8 @@ public void testTaskSubmission() throws Exception {
  CompletableFuture<Boolean> completionFuture = TestInvokable.COMPLETABLE_FUTURE;
 
  completionFuture.get();
+
+ taskInTerminalState.await();
  } finally {
  RpcUtils.terminateRpcEndpoint(taskManager, timeout);
  }