@Test
  public void testAbortJobCalledAfterKillingTasks() throws IOException {
    Configuration conf = new Configuration();
    conf.set(MRJobConfig.MR_AM_STAGING_DIR, stagingDir);
    conf.set(MRJobConfig.MR_AM_COMMITTER_CANCEL_TIMEOUT_MS, "1000");
    InlineDispatcher dispatcher = new InlineDispatcher();
    dispatcher.init(conf);
    dispatcher.start();
    OutputCommitter committer = Mockito.mock(OutputCommitter.class);
    CommitterEventHandler commitHandler = createCommitterEventHandler(dispatcher, committer);
    commitHandler.init(conf);
    commitHandler.start();
    JobImpl job = createRunningStubbedJob(conf, dispatcher, 2, null);

    // Fail one task. This should land the JobImpl in the FAIL_WAIT state
    job.handle(
        new JobTaskEvent(MRBuilderUtils.newTaskId(job.getID(), 1, TaskType.MAP), TaskState.FAILED));
    // Verify abort job hasn't been called
    Mockito.verify(committer, Mockito.never())
        .abortJob((JobContext) Mockito.any(), (State) Mockito.any());
    assertJobState(job, JobStateInternal.FAIL_WAIT);

    // Verify abortJob is called once and the job failed
    Mockito.verify(committer, Mockito.timeout(2000).times(1))
        .abortJob((JobContext) Mockito.any(), (State) Mockito.any());
    assertJobState(job, JobStateInternal.FAILED);

    dispatcher.stop();
  }
  @Before
  public void setUp() throws Exception {
    InlineDispatcher rmDispatcher = new InlineDispatcher();

    YarnConfiguration conf = new YarnConfiguration();
    YarnAPIStorageFactory.setConfiguration(conf);
    RMStorageFactory.setConfiguration(conf);
    conf.setClass(YarnConfiguration.RM_SCHEDULER, FifoScheduler.class, ResourceScheduler.class);
    RMUtilities.InitializeDB();
    TransactionStateManager tsm = new TransactionStateManager();
    tsm.init(conf);
    tsm.start();
    rmContext =
        new RMContextImpl(
            rmDispatcher,
            null,
            null,
            null,
            mock(DelegationTokenRenewer.class),
            null,
            null,
            null,
            conf,
            tsm);
    NodesListManager nodesListManager = mock(NodesListManager.class);
    HostsFileReader reader = mock(HostsFileReader.class);
    when(nodesListManager.getHostsReader()).thenReturn(reader);
    ((RMContextImpl) rmContext).setNodesListManager(nodesListManager);
    scheduler = mock(YarnScheduler.class);
    doAnswer(
            new Answer<Void>() {
              @Override
              public Void answer(InvocationOnMock invocation) throws Throwable {
                final SchedulerEvent event = (SchedulerEvent) (invocation.getArguments()[0]);
                eventType = event.getType();
                if (eventType == SchedulerEventType.NODE_UPDATE) {
                  List<UpdatedContainerInfo> lastestContainersInfoList =
                      ((NodeUpdateSchedulerEvent) event)
                          .getRMNode()
                          .pullContainerUpdates(event.getTransactionState());
                  for (UpdatedContainerInfo lastestContainersInfo : lastestContainersInfoList) {
                    completedContainers.addAll(lastestContainersInfo.getCompletedContainers());
                  }
                }
                return null;
              }
            })
        .when(scheduler)
        .handle(any(SchedulerEvent.class));

    rmDispatcher.register(SchedulerEventType.class, new TestSchedulerEventDispatcher());

    rmDispatcher.register(
        NodesListManagerEventType.class, new TestNodeListManagerEventDispatcher());

    NodeId nodeId = BuilderUtils.newNodeId("localhost", 0);
    node = new RMNodeImpl(nodeId, rmContext, nodeId.getHost(), 0, 0, null, null, null);
    nodesListManagerEvent = null;
  }