@Test
public void testCanceled() {
StageStateMachine stateMachine = createStageStateMachine();
assertTrue(stateMachine.transitionToCanceled());
assertFinalState(stateMachine, StageState.CANCELED);
}
@Test
public void testAborted() {
StageStateMachine stateMachine = createStageStateMachine();
assertTrue(stateMachine.transitionToAborted());
assertFinalState(stateMachine, StageState.ABORTED);
}
@Test
public void testFailed() {
StageStateMachine stateMachine = createStageStateMachine();
assertTrue(stateMachine.transitionToFailed(FAILED_CAUSE));
assertFinalState(stateMachine, StageState.FAILED);
}
@Test
public void testFinished() {
StageStateMachine stateMachine = createStageStateMachine();
assertTrue(stateMachine.transitionToFinished());
assertFinalState(stateMachine, StageState.FINISHED);
}
@Test
public void testBasicStateChanges() {
StageStateMachine stateMachine = createStageStateMachine();
assertState(stateMachine, StageState.PLANNED);
assertTrue(stateMachine.transitionToScheduling());
assertState(stateMachine, StageState.SCHEDULING);
assertTrue(stateMachine.transitionToScheduled());
assertState(stateMachine, StageState.SCHEDULED);
assertTrue(stateMachine.transitionToRunning());
assertState(stateMachine, StageState.RUNNING);
assertTrue(stateMachine.transitionToFinished());
assertState(stateMachine, StageState.FINISHED);
}
@Test
public void testPlanned() {
StageStateMachine stateMachine = createStageStateMachine();
assertState(stateMachine, StageState.PLANNED);
stateMachine = createStageStateMachine();
assertTrue(stateMachine.transitionToScheduling());
assertState(stateMachine, StageState.SCHEDULING);
stateMachine = createStageStateMachine();
assertTrue(stateMachine.transitionToRunning());
assertState(stateMachine, StageState.RUNNING);
stateMachine = createStageStateMachine();
assertTrue(stateMachine.transitionToFinished());
assertState(stateMachine, StageState.FINISHED);
stateMachine = createStageStateMachine();
assertTrue(stateMachine.transitionToFailed(FAILED_CAUSE));
assertState(stateMachine, StageState.FAILED);
stateMachine = createStageStateMachine();
assertTrue(stateMachine.transitionToAborted());
assertState(stateMachine, StageState.ABORTED);
stateMachine = createStageStateMachine();
assertTrue(stateMachine.transitionToCanceled());
assertState(stateMachine, StageState.CANCELED);
}
private static void assertFinalState(StageStateMachine stateMachine, StageState expectedState) {
assertTrue(expectedState.isDone());
assertState(stateMachine, expectedState);
assertFalse(stateMachine.transitionToScheduling());
assertState(stateMachine, expectedState);
assertFalse(stateMachine.transitionToScheduled());
assertState(stateMachine, expectedState);
assertFalse(stateMachine.transitionToRunning());
assertState(stateMachine, expectedState);
assertFalse(stateMachine.transitionToFinished());
assertState(stateMachine, expectedState);
assertFalse(stateMachine.transitionToFailed(FAILED_CAUSE));
assertState(stateMachine, expectedState);
assertFalse(stateMachine.transitionToAborted());
assertState(stateMachine, expectedState);
// attempt to fail with another exception, which will fail
assertFalse(stateMachine.transitionToFailed(new IOException("failure after finish")));
assertState(stateMachine, expectedState);
}
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());
}
}
