Example usage for org.apache.hadoop.yarn.api.records NodeId newInstance

List of usage examples for org.apache.hadoop.yarn.api.records NodeId newInstance

Introduction

In this page you can find the example usage for org.apache.hadoop.yarn.api.records NodeId newInstance.

Prototype

@Public
    @Stable
    public static NodeId newInstance(String host, int port) 

Source Link

Usage

From source file:org.apache.ignite.yarn.IgniteApplicationMasterSelfTest.java

License:Apache License

/**
 * @param host Host.//from   ww  w . j  av  a 2  s.c  o m
 * @param cpu Cpu count.
 * @param mem Memory.
 * @return Container.
 */
private Container createContainer(String host, int cpu, int mem) {
    return Container.newInstance(
            ContainerId.newContainerId(ApplicationAttemptId.newInstance(ApplicationId.newInstance(0l, 0), 0),
                    ThreadLocalRandom.current().nextLong()),
            NodeId.newInstance(host, 0), "example.com", new MockResource(mem, cpu), Priority.newInstance(0),
            null);
}

From source file:org.apache.myriad.TestObjectFactory.java

License:Apache License

public static RMNode getRMNode(String host, int port, Resource resource) {
    NodeId id = NodeId.newInstance(host, port);
    RMContext context = new MockRMContext();
    return new RMNodeImpl(id, context, id.getHost(), id.getPort(), id.getPort(), new NodeBase(host, "/tmp"),
            resource, "version-one");
}

From source file:org.apache.samza.job.yarn.TestYarnClusterResourceManager.java

License:Apache License

@Test
public void testErrorInStartContainerShouldUpdateState() {
    // create mocks
    final int samzaContainerId = 1;
    YarnConfiguration yarnConfiguration = mock(YarnConfiguration.class);
    SamzaAppMasterMetrics metrics = mock(SamzaAppMasterMetrics.class);
    Config config = mock(Config.class);
    AMRMClientAsync asyncClient = mock(AMRMClientAsync.class);
    YarnAppState yarnAppState = new YarnAppState(0, mock(ContainerId.class), "host", 8080, 8081);
    SamzaYarnAppMasterLifecycle lifecycle = mock(SamzaYarnAppMasterLifecycle.class);
    SamzaYarnAppMasterService service = mock(SamzaYarnAppMasterService.class);
    NMClientAsync asyncNMClient = mock(NMClientAsync.class);
    ClusterResourceManager.Callback callback = mock(ClusterResourceManager.Callback.class);

    // start the cluster manager
    YarnClusterResourceManager yarnClusterResourceManager = new YarnClusterResourceManager(asyncClient,
            asyncNMClient, callback, yarnAppState, lifecycle, service, metrics, yarnConfiguration, config);

    yarnAppState.pendingProcessors.put(String.valueOf(samzaContainerId),
            new YarnContainer(Container.newInstance(
                    ContainerId.newContainerId(
                            ApplicationAttemptId.newInstance(ApplicationId.newInstance(10000l, 1), 1), 1),
                    NodeId.newInstance("host1", 8088), "http://host1", Resource.newInstance(1024, 1),
                    Priority.newInstance(1),
                    Token.newInstance("id".getBytes(), "read", "password".getBytes(), "service"))));

    yarnClusterResourceManager.start();/* www .ja  v  a2s  .  co m*/
    assertEquals(1, yarnAppState.pendingProcessors.size());

    yarnClusterResourceManager.onStartContainerError(
            ContainerId.newContainerId(
                    ApplicationAttemptId.newInstance(ApplicationId.newInstance(10000l, 1), 1), 1),
            new Exception());

    assertEquals(0, yarnAppState.pendingProcessors.size());
    verify(callback, times(1)).onStreamProcessorLaunchFailure(anyObject(), any(Exception.class));
}

From source file:org.apache.samza.job.yarn.util.TestUtil.java

License:Apache License

public static Container getContainer(final ContainerId containerId, final String host, final int port) {
    return new Container() {
        @Override//from   www.  ja va 2  s .  co m
        public ContainerId getId() {
            return containerId;
        }

        @Override
        public void setId(ContainerId id) {
        }

        @Override
        public NodeId getNodeId() {
            return NodeId.newInstance(host, port);
        }

        @Override
        public void setNodeId(NodeId nodeId) {
        }

        @Override
        public String getNodeHttpAddress() {
            return host + ":" + port;
        }

        @Override
        public void setNodeHttpAddress(String nodeHttpAddress) {
        }

        @Override
        public Resource getResource() {
            return null;
        }

        @Override
        public void setResource(Resource resource) {
        }

        @Override
        public Priority getPriority() {
            return null;
        }

        @Override
        public void setPriority(Priority priority) {
        }

        @Override
        public Token getContainerToken() {
            return null;
        }

        @Override
        public void setContainerToken(Token containerToken) {
        }

        @Override
        public int compareTo(Container o) {
            return containerId.compareTo(o.getId());
        }
    };
}

From source file:org.apache.slider.server.appmaster.state.AppState.java

License:Apache License

/**
 * build up the special master node, which lives
 * in the live node set but has a lifecycle bonded to the AM
 * @param containerId the AM master/*from  w ww . j  a v a 2s .c  om*/
 * @param host hostname
 * @param amPort port
 * @param nodeHttpAddress http address: may be null
 */
public void buildAppMasterNode(ContainerId containerId, String host, int amPort, String nodeHttpAddress) {
    Container container = new ContainerPBImpl();
    container.setId(containerId);
    NodeId nodeId = NodeId.newInstance(host, amPort);
    container.setNodeId(nodeId);
    container.setNodeHttpAddress(nodeHttpAddress);
    RoleInstance am = new RoleInstance(container);
    am.role = SliderKeys.COMPONENT_AM;
    appMasterNode = am;
    //it is also added to the set of live nodes
    getLiveNodes().put(containerId, am);

    // patch up the role status
    RoleStatus roleStatus = roleStatusMap.get((SliderKeys.ROLE_AM_PRIORITY_INDEX));
    roleStatus.setDesired(1);
    roleStatus.incActual();
    roleStatus.incStarted();
}

From source file:org.apache.tez.dag.app.dag.impl.TestTaskAttempt.java

License:Apache License

@Test(timeout = 5000)
// Ensure the dag does not go into an error state if a attempt kill is
// received while STARTING
public void testLaunchFailedWhileKilling() throws Exception {
    ApplicationId appId = ApplicationId.newInstance(1, 2);
    ApplicationAttemptId appAttemptId = ApplicationAttemptId.newInstance(appId, 0);
    TezDAGID dagID = TezDAGID.getInstance(appId, 1);
    TezVertexID vertexID = TezVertexID.getInstance(dagID, 1);
    TezTaskID taskID = TezTaskID.getInstance(vertexID, 1);
    TezTaskAttemptID taskAttemptID = TezTaskAttemptID.getInstance(taskID, 0);

    MockEventHandler eventHandler = new MockEventHandler();
    TaskAttemptListener taListener = mock(TaskAttemptListener.class);
    when(taListener.getAddress()).thenReturn(new InetSocketAddress("localhost", 0));

    Configuration taskConf = new Configuration();
    taskConf.setClass("fs.file.impl", StubbedFS.class, FileSystem.class);
    taskConf.setBoolean("fs.file.impl.disable.cache", true);

    TaskLocationHint locationHint = TaskLocationHint
            .createTaskLocationHint(new HashSet<String>(Arrays.asList(new String[] { "127.0.0.1" })), null);
    Resource resource = Resource.newInstance(1024, 1);

    AppContext mockAppContext = mock(AppContext.class);
    doReturn(new ClusterInfo()).when(mockAppContext).getClusterInfo();

    TaskAttemptImpl taImpl = new MockTaskAttemptImpl(taskID, 1, eventHandler, taListener, taskConf,
            new SystemClock(), mock(TaskHeartbeatHandler.class), mockAppContext, locationHint, false, resource,
            createFakeContainerContext(), false);

    NodeId nid = NodeId.newInstance("127.0.0.1", 0);
    ContainerId contId = ContainerId.newInstance(appAttemptId, 3);
    Container container = mock(Container.class);
    when(container.getId()).thenReturn(contId);
    when(container.getNodeId()).thenReturn(nid);

    taImpl.handle(new TaskAttemptEventSchedule(taskAttemptID, 0, 0));
    // At state STARTING.
    taImpl.handle(new TaskAttemptEventKillRequest(taskAttemptID, null,
            TaskAttemptTerminationCause.TERMINATED_BY_CLIENT));
    // At some KILLING state.
    taImpl.handle(new TaskAttemptEventKillRequest(taskAttemptID, null,
            TaskAttemptTerminationCause.TERMINATED_BY_CLIENT));
    // taImpl.handle(new TaskAttemptEventContainerTerminating(taskAttemptID,
    // null));/*from w  w  w . ja  va  2  s  .  co m*/
    assertFalse(eventHandler.internalError);
}

From source file:org.apache.tez.dag.app.dag.impl.TestTaskAttempt.java

License:Apache License

@Test(timeout = 5000)
// Ensure ContainerTerminating and ContainerTerminated is handled correctly by
// the TaskAttempt
public void testContainerTerminationWhileRunning() throws Exception {
    ApplicationId appId = ApplicationId.newInstance(1, 2);
    ApplicationAttemptId appAttemptId = ApplicationAttemptId.newInstance(appId, 0);
    TezDAGID dagID = TezDAGID.getInstance(appId, 1);
    TezVertexID vertexID = TezVertexID.getInstance(dagID, 1);
    TezTaskID taskID = TezTaskID.getInstance(vertexID, 1);

    MockEventHandler eventHandler = spy(new MockEventHandler());
    TaskAttemptListener taListener = mock(TaskAttemptListener.class);
    when(taListener.getAddress()).thenReturn(new InetSocketAddress("localhost", 0));

    Configuration taskConf = new Configuration();
    taskConf.setClass("fs.file.impl", StubbedFS.class, FileSystem.class);
    taskConf.setBoolean("fs.file.impl.disable.cache", true);

    TaskLocationHint locationHint = TaskLocationHint
            .createTaskLocationHint(new HashSet<String>(Arrays.asList(new String[] { "127.0.0.1" })), null);
    Resource resource = Resource.newInstance(1024, 1);

    NodeId nid = NodeId.newInstance("127.0.0.1", 0);
    ContainerId contId = ContainerId.newInstance(appAttemptId, 3);
    Container container = mock(Container.class);
    when(container.getId()).thenReturn(contId);
    when(container.getNodeId()).thenReturn(nid);
    when(container.getNodeHttpAddress()).thenReturn("localhost:0");

    AppContext appCtx = mock(AppContext.class);
    AMContainerMap containers = new AMContainerMap(mock(ContainerHeartbeatHandler.class),
            mock(TaskAttemptListener.class), new ContainerContextMatcher(), appCtx);
    containers.addContainerIfNew(container);

    doReturn(new ClusterInfo()).when(appCtx).getClusterInfo();
    doReturn(containers).when(appCtx).getAllContainers();

    TaskHeartbeatHandler mockHeartbeatHandler = mock(TaskHeartbeatHandler.class);
    TaskAttemptImpl taImpl = new MockTaskAttemptImpl(taskID, 1, eventHandler, taListener, taskConf,
            new SystemClock(), mockHeartbeatHandler, appCtx, locationHint, false, resource,
            createFakeContainerContext(), false);
    TezTaskAttemptID taskAttemptID = taImpl.getID();
    ArgumentCaptor<Event> arg = ArgumentCaptor.forClass(Event.class);

    taImpl.handle(new TaskAttemptEventSchedule(taskAttemptID, 0, 0));
    // At state STARTING.
    taImpl.handle(new TaskAttemptEventStartedRemotely(taskAttemptID, contId, null));
    assertEquals("Task attempt is not in the RUNNING state", taImpl.getState(), TaskAttemptState.RUNNING);
    verify(mockHeartbeatHandler).register(taskAttemptID);

    int expectedEventsAtRunning = 3;
    verify(eventHandler, times(expectedEventsAtRunning)).handle(arg.capture());

    taImpl.handle(new TaskAttemptEventContainerTerminating(taskAttemptID, "Terminating",
            TaskAttemptTerminationCause.APPLICATION_ERROR));
    assertFalse("InternalError occurred trying to handle TA_CONTAINER_TERMINATING", eventHandler.internalError);
    verify(mockHeartbeatHandler).unregister(taskAttemptID);
    assertEquals("Task attempt is not in the  FAILED state", taImpl.getState(), TaskAttemptState.FAILED);

    assertEquals(1, taImpl.getDiagnostics().size());
    assertEquals("Terminating", taImpl.getDiagnostics().get(0));
    assertEquals(TaskAttemptTerminationCause.APPLICATION_ERROR, taImpl.getTerminationCause());

    int expectedEvenstAfterTerminating = expectedEventsAtRunning + 3;
    arg = ArgumentCaptor.forClass(Event.class);
    verify(eventHandler, times(expectedEvenstAfterTerminating)).handle(arg.capture());

    verifyEventType(arg.getAllValues().subList(expectedEventsAtRunning, expectedEvenstAfterTerminating),
            TaskEventTAUpdate.class, 1);
    verifyEventType(arg.getAllValues().subList(expectedEventsAtRunning, expectedEvenstAfterTerminating),
            AMSchedulerEventTAEnded.class, 1);
    verifyEventType(arg.getAllValues().subList(expectedEventsAtRunning, expectedEvenstAfterTerminating),
            DAGEventCounterUpdate.class, 1);

    taImpl.handle(new TaskAttemptEventContainerTerminated(taskAttemptID, "Terminated",
            TaskAttemptTerminationCause.CONTAINER_EXITED));
    // verify unregister is not invoked again
    verify(mockHeartbeatHandler, times(1)).unregister(taskAttemptID);
    int expectedEventAfterTerminated = expectedEvenstAfterTerminating + 0;
    arg = ArgumentCaptor.forClass(Event.class);
    verify(eventHandler, times(expectedEventAfterTerminated)).handle(arg.capture());

    assertEquals(2, taImpl.getDiagnostics().size());
    assertEquals("Terminated", taImpl.getDiagnostics().get(1));

    // check that original error cause is retained
    assertEquals(TaskAttemptTerminationCause.APPLICATION_ERROR, taImpl.getTerminationCause());
}

From source file:org.apache.tez.dag.app.dag.impl.TestTaskAttempt.java

License:Apache License

@Test(timeout = 5000)
// Ensure ContainerTerminated is handled correctly by the TaskAttempt
public void testContainerTerminatedWhileRunning() throws Exception {
    ApplicationId appId = ApplicationId.newInstance(1, 2);
    ApplicationAttemptId appAttemptId = ApplicationAttemptId.newInstance(appId, 0);
    TezDAGID dagID = TezDAGID.getInstance(appId, 1);
    TezVertexID vertexID = TezVertexID.getInstance(dagID, 1);
    TezTaskID taskID = TezTaskID.getInstance(vertexID, 1);

    MockEventHandler eventHandler = new MockEventHandler();
    TaskAttemptListener taListener = mock(TaskAttemptListener.class);
    when(taListener.getAddress()).thenReturn(new InetSocketAddress("localhost", 0));

    Configuration taskConf = new Configuration();
    taskConf.setClass("fs.file.impl", StubbedFS.class, FileSystem.class);
    taskConf.setBoolean("fs.file.impl.disable.cache", true);

    TaskLocationHint locationHint = TaskLocationHint
            .createTaskLocationHint(new HashSet<String>(Arrays.asList(new String[] { "127.0.0.1" })), null);
    Resource resource = Resource.newInstance(1024, 1);

    NodeId nid = NodeId.newInstance("127.0.0.1", 0);
    ContainerId contId = ContainerId.newInstance(appAttemptId, 3);
    Container container = mock(Container.class);
    when(container.getId()).thenReturn(contId);
    when(container.getNodeId()).thenReturn(nid);
    when(container.getNodeHttpAddress()).thenReturn("localhost:0");

    AppContext appCtx = mock(AppContext.class);
    AMContainerMap containers = new AMContainerMap(mock(ContainerHeartbeatHandler.class),
            mock(TaskAttemptListener.class), new ContainerContextMatcher(), appCtx);
    containers.addContainerIfNew(container);

    doReturn(new ClusterInfo()).when(appCtx).getClusterInfo();
    doReturn(containers).when(appCtx).getAllContainers();

    TaskHeartbeatHandler mockHeartbeatHandler = mock(TaskHeartbeatHandler.class);
    TaskAttemptImpl taImpl = new MockTaskAttemptImpl(taskID, 1, eventHandler, taListener, taskConf,
            new SystemClock(), mockHeartbeatHandler, appCtx, locationHint, false, resource,
            createFakeContainerContext(), false);
    TezTaskAttemptID taskAttemptID = taImpl.getID();
    taImpl.handle(new TaskAttemptEventSchedule(taskAttemptID, 0, 0));
    // At state STARTING.
    taImpl.handle(new TaskAttemptEventStartedRemotely(taskAttemptID, contId, null));
    assertEquals("Task attempt is not in running state", taImpl.getState(), TaskAttemptState.RUNNING);
    verify(mockHeartbeatHandler).register(taskAttemptID);

    taImpl.handle(new TaskAttemptEventContainerTerminated(taskAttemptID, "Terminated",
            TaskAttemptTerminationCause.CONTAINER_EXITED));
    assertFalse("InternalError occurred trying to handle TA_CONTAINER_TERMINATED", eventHandler.internalError);
    verify(mockHeartbeatHandler).unregister(taskAttemptID);
    assertEquals("Terminated", taImpl.getDiagnostics().get(0));
    assertEquals(TaskAttemptTerminationCause.CONTAINER_EXITED, taImpl.getTerminationCause());
    // TODO Ensure TA_TERMINATING after this is ingored.
}

From source file:org.apache.tez.dag.app.dag.impl.TestTaskAttempt.java

License:Apache License

@Test(timeout = 5000)
// Ensure ContainerTerminating and ContainerTerminated is handled correctly by
// the TaskAttempt
public void testContainerTerminatedAfterSuccess() throws Exception {
    ApplicationId appId = ApplicationId.newInstance(1, 2);
    ApplicationAttemptId appAttemptId = ApplicationAttemptId.newInstance(appId, 0);
    TezDAGID dagID = TezDAGID.getInstance(appId, 1);
    TezVertexID vertexID = TezVertexID.getInstance(dagID, 1);
    TezTaskID taskID = TezTaskID.getInstance(vertexID, 1);

    MockEventHandler eventHandler = spy(new MockEventHandler());
    TaskAttemptListener taListener = mock(TaskAttemptListener.class);
    when(taListener.getAddress()).thenReturn(new InetSocketAddress("localhost", 0));

    Configuration taskConf = new Configuration();
    taskConf.setClass("fs.file.impl", StubbedFS.class, FileSystem.class);
    taskConf.setBoolean("fs.file.impl.disable.cache", true);

    TaskLocationHint locationHint = TaskLocationHint
            .createTaskLocationHint(new HashSet<String>(Arrays.asList(new String[] { "127.0.0.1" })), null);
    Resource resource = Resource.newInstance(1024, 1);

    NodeId nid = NodeId.newInstance("127.0.0.1", 0);
    ContainerId contId = ContainerId.newInstance(appAttemptId, 3);
    Container container = mock(Container.class);
    when(container.getId()).thenReturn(contId);
    when(container.getNodeId()).thenReturn(nid);
    when(container.getNodeHttpAddress()).thenReturn("localhost:0");

    AppContext appCtx = mock(AppContext.class);
    AMContainerMap containers = new AMContainerMap(mock(ContainerHeartbeatHandler.class),
            mock(TaskAttemptListener.class), new ContainerContextMatcher(), appCtx);
    containers.addContainerIfNew(container);

    doReturn(new ClusterInfo()).when(appCtx).getClusterInfo();
    doReturn(containers).when(appCtx).getAllContainers();

    TaskHeartbeatHandler mockHeartbeatHandler = mock(TaskHeartbeatHandler.class);
    TaskAttemptImpl taImpl = new MockTaskAttemptImpl(taskID, 1, eventHandler, taListener, taskConf,
            new SystemClock(), mockHeartbeatHandler, appCtx, locationHint, false, resource,
            createFakeContainerContext(), false);
    TezTaskAttemptID taskAttemptID = taImpl.getID();
    ArgumentCaptor<Event> arg = ArgumentCaptor.forClass(Event.class);

    taImpl.handle(new TaskAttemptEventSchedule(taskAttemptID, 0, 0));
    // At state STARTING.
    taImpl.handle(new TaskAttemptEventStartedRemotely(taskAttemptID, contId, null));
    assertEquals("Task attempt is not in the RUNNING state", taImpl.getState(), TaskAttemptState.RUNNING);
    verify(mockHeartbeatHandler).register(taskAttemptID);

    int expectedEventsAtRunning = 3;
    verify(eventHandler, times(expectedEventsAtRunning)).handle(arg.capture());

    taImpl.handle(new TaskAttemptEvent(taskAttemptID, TaskAttemptEventType.TA_DONE));

    assertEquals("Task attempt is not in the  SUCCEEDED state", taImpl.getState(), TaskAttemptState.SUCCEEDED);
    verify(mockHeartbeatHandler).unregister(taskAttemptID);
    assertEquals(0, taImpl.getDiagnostics().size());

    int expectedEvenstAfterTerminating = expectedEventsAtRunning + 3;
    arg = ArgumentCaptor.forClass(Event.class);
    verify(eventHandler, times(expectedEvenstAfterTerminating)).handle(arg.capture());

    verifyEventType(arg.getAllValues().subList(expectedEventsAtRunning, expectedEvenstAfterTerminating),
            TaskEventTAUpdate.class, 1);
    verifyEventType(arg.getAllValues().subList(expectedEventsAtRunning, expectedEvenstAfterTerminating),
            AMSchedulerEventTAEnded.class, 1);
    verifyEventType(arg.getAllValues().subList(expectedEventsAtRunning, expectedEvenstAfterTerminating),
            DAGEventCounterUpdate.class, 1);

    taImpl.handle(new TaskAttemptEventContainerTerminated(taskAttemptID, "Terminated",
            TaskAttemptTerminationCause.CONTAINER_EXITED));
    // verify unregister is not invoked again
    verify(mockHeartbeatHandler, times(1)).unregister(taskAttemptID);
    int expectedEventAfterTerminated = expectedEvenstAfterTerminating + 0;
    arg = ArgumentCaptor.forClass(Event.class);
    verify(eventHandler, times(expectedEventAfterTerminated)).handle(arg.capture());

    // Verify that the diagnostic message included in the Terminated event is not
    // captured - TA already succeeded. Error cause is the default value.
    assertEquals(0, taImpl.getDiagnostics().size());
    assertEquals(TaskAttemptTerminationCause.UNKNOWN_ERROR, taImpl.getTerminationCause());
}

From source file:org.apache.tez.dag.app.dag.impl.TestTaskAttempt.java

License:Apache License

@Test(timeout = 5000)
public void testFailure() throws Exception {
    ApplicationId appId = ApplicationId.newInstance(1, 2);
    ApplicationAttemptId appAttemptId = ApplicationAttemptId.newInstance(appId, 0);
    TezDAGID dagID = TezDAGID.getInstance(appId, 1);
    TezVertexID vertexID = TezVertexID.getInstance(dagID, 1);
    TezTaskID taskID = TezTaskID.getInstance(vertexID, 1);

    MockEventHandler eventHandler = spy(new MockEventHandler());
    TaskAttemptListener taListener = mock(TaskAttemptListener.class);
    when(taListener.getAddress()).thenReturn(new InetSocketAddress("localhost", 0));

    Configuration taskConf = new Configuration();
    taskConf.setClass("fs.file.impl", StubbedFS.class, FileSystem.class);
    taskConf.setBoolean("fs.file.impl.disable.cache", true);
    taskConf.setBoolean(TezConfiguration.TEZ_AM_SPECULATION_ENABLED, true);

    TaskLocationHint locationHint = TaskLocationHint
            .createTaskLocationHint(new HashSet<String>(Arrays.asList(new String[] { "127.0.0.1" })), null);
    Resource resource = Resource.newInstance(1024, 1);

    NodeId nid = NodeId.newInstance("127.0.0.1", 0);
    ContainerId contId = ContainerId.newInstance(appAttemptId, 3);
    Container container = mock(Container.class);
    when(container.getId()).thenReturn(contId);
    when(container.getNodeId()).thenReturn(nid);
    when(container.getNodeHttpAddress()).thenReturn("localhost:0");

    AppContext appCtx = mock(AppContext.class);
    AMContainerMap containers = new AMContainerMap(mock(ContainerHeartbeatHandler.class),
            mock(TaskAttemptListener.class), new ContainerContextMatcher(), appCtx);
    containers.addContainerIfNew(container);

    doReturn(new ClusterInfo()).when(appCtx).getClusterInfo();
    doReturn(containers).when(appCtx).getAllContainers();

    TaskHeartbeatHandler mockHeartbeatHandler = mock(TaskHeartbeatHandler.class);
    TaskAttemptImpl taImpl = new MockTaskAttemptImpl(taskID, 1, eventHandler, taListener, taskConf,
            new SystemClock(), mockHeartbeatHandler, appCtx, locationHint, false, resource,
            createFakeContainerContext(), false);
    TezTaskAttemptID taskAttemptID = taImpl.getID();
    ArgumentCaptor<Event> arg = ArgumentCaptor.forClass(Event.class);

    taImpl.handle(new TaskAttemptEventSchedule(taskAttemptID, 0, 0));
    // At state STARTING.
    taImpl.handle(new TaskAttemptEventStartedRemotely(taskAttemptID, contId, null));
    assertEquals("Task attempt is not in the RUNNING state", taImpl.getState(), TaskAttemptState.RUNNING);
    verify(mockHeartbeatHandler).register(taskAttemptID);

    int expectedEventsAtRunning = 4;
    verify(eventHandler, times(expectedEventsAtRunning)).handle(arg.capture());
    verifyEventType(arg.getAllValues().subList(0, expectedEventsAtRunning),
            SpeculatorEventTaskAttemptStatusUpdate.class, 1);

    taImpl.handle(new TaskAttemptEventStatusUpdate(taskAttemptID, new TaskStatusUpdateEvent(null, 0.1f)));

    taImpl.handle(new TaskAttemptEventAttemptFailed(taskAttemptID, TaskAttemptEventType.TA_FAILED, "0",
            TaskAttemptTerminationCause.APPLICATION_ERROR));

    assertEquals("Task attempt is not in the  FAIL_IN_PROGRESS state", taImpl.getInternalState(),
            TaskAttemptStateInternal.FAIL_IN_PROGRESS);
    verify(mockHeartbeatHandler).unregister(taskAttemptID);
    assertEquals(1, taImpl.getDiagnostics().size());
    assertEquals("0", taImpl.getDiagnostics().get(0));
    assertEquals(TaskAttemptTerminationCause.APPLICATION_ERROR, taImpl.getTerminationCause());

    taImpl.handle(new TaskAttemptEventContainerTerminated(taskAttemptID, "1",
            TaskAttemptTerminationCause.CONTAINER_EXITED));
    // verify unregister is not invoked again
    verify(mockHeartbeatHandler, times(1)).unregister(taskAttemptID);
    assertEquals(2, taImpl.getDiagnostics().size());
    assertEquals("1", taImpl.getDiagnostics().get(1));
    // err cause does not change
    assertEquals(TaskAttemptTerminationCause.APPLICATION_ERROR, taImpl.getTerminationCause());

    int expectedEvenstAfterTerminating = expectedEventsAtRunning + 5;
    arg = ArgumentCaptor.forClass(Event.class);
    verify(eventHandler, times(expectedEvenstAfterTerminating)).handle(arg.capture());

    Event e = verifyEventType(/*  w w w . j  a  va  2s .c  o  m*/
            arg.getAllValues().subList(expectedEventsAtRunning, expectedEvenstAfterTerminating),
            TaskEventTAUpdate.class, 1);
    assertEquals(TaskEventType.T_ATTEMPT_FAILED, e.getType());
    verifyEventType(arg.getAllValues().subList(expectedEventsAtRunning, expectedEvenstAfterTerminating),
            AMSchedulerEventTAEnded.class, 1);
    verifyEventType(arg.getAllValues().subList(expectedEventsAtRunning, expectedEvenstAfterTerminating),
            DAGEventCounterUpdate.class, 1);
    verifyEventType(arg.getAllValues().subList(expectedEventsAtRunning, expectedEvenstAfterTerminating),
            SpeculatorEventTaskAttemptStatusUpdate.class, 2);
}