private static void assertState(StageStateMachine stateMachine, StageState expectedState) {
assertEquals(stateMachine.getStageId(), STAGE_ID);
assertEquals(stateMachine.getLocation(), LOCATION);
assertSame(stateMachine.getSession(), TEST_SESSION);
StageInfo stageInfo = stateMachine.getStageInfo(ImmutableList::of, ImmutableList::of);
assertEquals(stageInfo.getStageId(), STAGE_ID);
assertEquals(stageInfo.getSelf(), LOCATION);
assertEquals(stageInfo.getSubStages(), ImmutableList.of());
assertEquals(stageInfo.getTasks(), ImmutableList.of());
assertEquals(stageInfo.getTypes(), ImmutableList.of(VARCHAR));
assertSame(stageInfo.getPlan(), PLAN_FRAGMENT);
assertEquals(stateMachine.getState(), expectedState);
assertEquals(stageInfo.getState(), expectedState);
if (expectedState == StageState.FAILED) {
ExecutionFailureInfo failure = stageInfo.getFailureCause();
assertEquals(failure.getMessage(), FAILED_CAUSE.getMessage());
assertEquals(failure.getType(), FAILED_CAUSE.getClass().getName());
} else {
assertNull(stageInfo.getFailureCause());
}
}
