@Test(timeout = 4000)
public void testWaitUntilFinish() throws Exception {

  final PhysicalPlan physicalPlan =
    TestPlanGenerator.generatePhysicalPlan(TestPlanGenerator.PlanType.TwoVerticesJoined, false);
  planStateManager.updatePlan(physicalPlan, MAX_SCHEDULE_ATTEMPT);
  assertFalse(planStateManager.isPlanDone());
  // Wait for the plan to finish and check the plan state.
  // It have to return EXECUTING state after timeout.
  final PlanState.State executingState = planStateManager.waitUntilFinish(100, TimeUnit.MILLISECONDS);
  assertEquals(PlanState.State.EXECUTING, executingState);
  // Complete the plan and check the result again.
  // It has to return COMPLETE.
  final List<String> tasks = physicalPlan.getStageDAG().getTopologicalSort().stream()
    .flatMap(stage -> planStateManager.getTaskAttemptsToSchedule(stage.getId()).stream())
    .collect(Collectors.toList());
  tasks.forEach(taskId -> planStateManager.onTaskStateChanged(taskId, TaskState.State.EXECUTING));
  tasks.forEach(taskId -> planStateManager.onTaskStateChanged(taskId, TaskState.State.COMPLETE));
  final PlanState.State completedState = planStateManager.waitUntilFinish();
  assertEquals(PlanState.State.COMPLETE, completedState);
}