Examples with ContainerHeartbeatHandler org.apache.tez.dag.app.ContainerHeartbeatHandler used on opensource projects

Example 1 with ContainerHeartbeatHandler

use of org.apache.tez.dag.app.ContainerHeartbeatHandler in project tez by apache.

the class TestAMContainerMap method testCleanupOnDagComplete.

@Test(timeout = 10000)
public void testCleanupOnDagComplete() {
    ContainerHeartbeatHandler chh = mock(ContainerHeartbeatHandler.class);
    TaskCommunicatorManagerInterface tal = mock(TaskCommunicatorManagerInterface.class);
    AppContext appContext = mock(AppContext.class);
    when(appContext.getAMConf()).thenReturn(new Configuration());
    int numContainers = 7;
    WrappedContainer[] wContainers = new WrappedContainer[numContainers];
    for (int i = 0; i < numContainers; i++) {
        WrappedContainer wc = new WrappedContainer(false, null, i);
        wContainers[i] = wc;
    }
    AMContainerMap amContainerMap = new AMContainerMapForTest(chh, tal, mock(ContainerSignatureMatcher.class), appContext, wContainers);
    for (int i = 0; i < numContainers; i++) {
        amContainerMap.addContainerIfNew(wContainers[i].container, 0, 0, 0);
    }
    // Container 1 in LAUNCHING state
    wContainers[0].launchContainer();
    wContainers[0].verifyState(AMContainerState.LAUNCHING);
    // Container 2 in IDLE state
    wContainers[1].launchContainer();
    wContainers[1].containerLaunched();
    wContainers[1].verifyState(AMContainerState.IDLE);
    // Container 3 RUNNING state
    wContainers[2].launchContainer();
    wContainers[2].containerLaunched();
    wContainers[2].assignTaskAttempt(wContainers[2].taskAttemptID);
    wContainers[2].verifyState(AMContainerState.RUNNING);
    // Cointainer 4 STOP_REQUESTED
    wContainers[3].launchContainer();
    wContainers[3].containerLaunched();
    wContainers[3].stopRequest();
    wContainers[3].verifyState(AMContainerState.STOP_REQUESTED);
    // Container 5 STOPPING
    wContainers[4].launchContainer();
    wContainers[4].containerLaunched();
    wContainers[4].stopRequest();
    wContainers[4].nmStopSent();
    wContainers[4].verifyState(AMContainerState.STOPPING);
    // Container 6 COMPLETED
    wContainers[5].launchContainer();
    wContainers[5].containerLaunched();
    wContainers[5].stopRequest();
    wContainers[5].nmStopSent();
    wContainers[5].containerCompleted();
    wContainers[5].verifyState(AMContainerState.COMPLETED);
    // Container 7 STOP_REQUESTED + ERROR
    wContainers[6].launchContainer();
    wContainers[6].containerLaunched();
    wContainers[6].containerLaunched();
    assertTrue(wContainers[6].amContainer.isInErrorState());
    wContainers[6].verifyState(AMContainerState.STOP_REQUESTED);
    // 7 containers present, and registered with AMContainerMap at this point.
    assertEquals(7, amContainerMap.containerMap.size());
    amContainerMap.dagComplete(mock(DAG.class));
    assertEquals(5, amContainerMap.containerMap.size());
}