VertexStatusBuilder mockVertexStatusBuilder =
mock(VertexStatusBuilder.class);
when(mockDAG.getVertexStatus(anyString(), anySetOf(StatusGetOpts.class)))
.thenReturn(mockVertexStatusBuilder);
DAGAppMaster mockDagAM = mock(DAGAppMaster.class);
AppContext mockAppContext = mock(AppContext.class);
when(mockDagAM.getContext()).thenReturn(mockAppContext);
when(mockDagAM.getContext().getCurrentDAG()).thenReturn(mockDAG);
DAGClientHandler dagClientHandler = new DAGClientHandler(mockDagAM);
// getAllDAGs()
assertEquals(1, dagClientHandler.getAllDAGs().size());
assertEquals("dag_9999_0001_1", dagClientHandler.getAllDAGs().get(0));
// getDAGStatus
try {
dagClientHandler.getDAGStatus("dag_9999_0001_2", Sets.newSet(StatusGetOpts.GET_COUNTERS));
fail("should not come here");
} catch (TezException e) {
assertTrue(e.getMessage().contains("Unknown dagId"));
}
DAGStatus dagStatus = dagClientHandler.getDAGStatus("dag_9999_0001_1",
Sets.newSet(StatusGetOpts.GET_COUNTERS));
assertEquals(mockDagStatusBuilder, dagStatus);
// getVertexStatus
try {
dagClientHandler.getVertexStatus("dag_9999_0001_2", "v1", Sets.newSet(StatusGetOpts.GET_COUNTERS));
fail("should not come here");
} catch (TezException e) {
assertTrue(e.getMessage().contains("Unknown dagId"));
}
VertexStatus vertexStatus = dagClientHandler.getVertexStatus("dag_9999_0001_1", "v1",
Sets.newSet(StatusGetOpts.GET_COUNTERS));
assertEquals(mockVertexStatusBuilder, vertexStatus);
// getSessionStatus
when(mockDagAM.isSession()).thenReturn(false);
try{
dagClientHandler.getSessionStatus();
fail("should not come here");
}catch(TezException e){
assertEquals("Unsupported operation as AM not running in session mode", e.getMessage());
}
when(mockDagAM.isSession()).thenReturn(true);
when(mockDagAM.getState()).thenReturn(DAGAppMasterState.INITED);
assertEquals(TezAppMasterStatus.INITIALIZING, dagClientHandler.getSessionStatus());
when(mockDagAM.getState()).thenReturn(DAGAppMasterState.ERROR);
assertEquals(TezAppMasterStatus.SHUTDOWN, dagClientHandler.getSessionStatus());
// tryKillDAG
try{
dagClientHandler.tryKillDAG("dag_9999_0001_2");