@Test(timeout = 5000)
  public void testExpiredContainer() {
    // Start the node
    node.handle(new RMNodeEvent(null, RMNodeEventType.STARTED, null));
    verify(scheduler).handle(any(NodeAddedSchedulerEvent.class));

    // Expire a container
    ContainerId completedContainerId =
        BuilderUtils.newContainerId(
            BuilderUtils.newApplicationAttemptId(BuilderUtils.newApplicationId(0, 0), 0), 0);
    node.handle(new RMNodeCleanContainerEvent(null, completedContainerId, null));
    Assert.assertEquals(1, node.getContainersToCleanUp().size());

    // Now verify that scheduler isn't notified of an expired container
    // by checking number of 'completedContainers' it got in the previous event
    RMNodeStatusEvent statusEvent = getMockRMNodeStatusEvent();
    ContainerStatus containerStatus = mock(ContainerStatus.class);
    doReturn(completedContainerId).when(containerStatus).getContainerId();
    doReturn(Collections.singletonList(containerStatus)).when(statusEvent).getContainers();
    node.handle(statusEvent);
    /* Expect the scheduler call handle function 2 times
     * 1. RMNode status from new to Running, handle the add_node event
     * 2. handle the node update event
     */
    verify(scheduler, times(2)).handle(any(NodeUpdateSchedulerEvent.class));
  }
  @Test(timeout = 500000)
  public void testStatusChange() {
    // Start the node
    node.handle(
        new RMNodeEvent(
            null, RMNodeEventType.STARTED, new TransactionStateImpl(TransactionType.RM)));
    // Add info to the queue first
    node.setNextHeartBeat(false);

    ContainerId completedContainerId1 =
        BuilderUtils.newContainerId(
            BuilderUtils.newApplicationAttemptId(BuilderUtils.newApplicationId(0, 0), 0), 0);
    ContainerId completedContainerId2 =
        BuilderUtils.newContainerId(
            BuilderUtils.newApplicationAttemptId(BuilderUtils.newApplicationId(1, 1), 1), 1);

    RMNodeStatusEvent statusEvent1 = getMockRMNodeStatusEvent();
    RMNodeStatusEvent statusEvent2 = getMockRMNodeStatusEvent();

    ContainerStatus containerStatus1 = mock(ContainerStatus.class);
    ContainerStatus containerStatus2 = mock(ContainerStatus.class);

    doReturn(completedContainerId1).when(containerStatus1).getContainerId();
    doReturn(Collections.singletonList(containerStatus1)).when(statusEvent1).getContainers();
    doReturn(ContainerState.COMPLETE).when(containerStatus1).getState();
    doReturn(completedContainerId2).when(containerStatus2).getContainerId();
    doReturn(Collections.singletonList(containerStatus2)).when(statusEvent2).getContainers();
    doReturn(ContainerState.COMPLETE).when(containerStatus2).getState();
    verify(scheduler, times(1)).handle(any(NodeUpdateSchedulerEvent.class));
    node.handle(statusEvent1);
    node.handle(statusEvent2);
    verify(scheduler, times(1)).handle(any(NodeUpdateSchedulerEvent.class));
    Assert.assertEquals(2, node.getQueueSize());
    node.handle(
        new RMNodeEvent(
            node.getNodeID(),
            RMNodeEventType.EXPIRE,
            new TransactionStateImpl(TransactionType.RM)));
    Assert.assertEquals(0, node.getQueueSize());
  }
  @Test
  public void testApplicationOrderingWithPriority() throws Exception {

    Configuration conf = new Configuration();
    conf.setClass(YarnConfiguration.RM_SCHEDULER, CapacityScheduler.class, ResourceScheduler.class);
    MockRM rm = new MockRM(conf);
    rm.start();
    CapacityScheduler cs = (CapacityScheduler) rm.getResourceScheduler();

    LeafQueue q = (LeafQueue) cs.getQueue("default");
    Assert.assertNotNull(q);

    String host = "127.0.0.1";
    RMNode node = MockNodes.newNodeInfo(0, MockNodes.newResource(16 * GB), 1, host);
    cs.handle(new NodeAddedSchedulerEvent(node));

    // add app 1 start
    ApplicationId appId1 = BuilderUtils.newApplicationId(100, 1);
    ApplicationAttemptId appAttemptId1 = BuilderUtils.newApplicationAttemptId(appId1, 1);

    RMAppAttemptMetrics attemptMetric1 = new RMAppAttemptMetrics(appAttemptId1, rm.getRMContext());
    RMAppImpl app1 = mock(RMAppImpl.class);
    when(app1.getApplicationId()).thenReturn(appId1);
    RMAppAttemptImpl attempt1 = mock(RMAppAttemptImpl.class);
    when(attempt1.getAppAttemptId()).thenReturn(appAttemptId1);
    when(attempt1.getRMAppAttemptMetrics()).thenReturn(attemptMetric1);
    when(app1.getCurrentAppAttempt()).thenReturn(attempt1);

    rm.getRMContext().getRMApps().put(appId1, app1);

    SchedulerEvent addAppEvent1 =
        new AppAddedSchedulerEvent(appId1, "default", "user", null, Priority.newInstance(5));
    cs.handle(addAppEvent1);
    SchedulerEvent addAttemptEvent1 = new AppAttemptAddedSchedulerEvent(appAttemptId1, false);
    cs.handle(addAttemptEvent1);
    // add app1 end

    // add app2 begin
    ApplicationId appId2 = BuilderUtils.newApplicationId(100, 2);
    ApplicationAttemptId appAttemptId2 = BuilderUtils.newApplicationAttemptId(appId2, 1);

    RMAppAttemptMetrics attemptMetric2 = new RMAppAttemptMetrics(appAttemptId2, rm.getRMContext());
    RMAppImpl app2 = mock(RMAppImpl.class);
    when(app2.getApplicationId()).thenReturn(appId2);
    RMAppAttemptImpl attempt2 = mock(RMAppAttemptImpl.class);
    when(attempt2.getAppAttemptId()).thenReturn(appAttemptId2);
    when(attempt2.getRMAppAttemptMetrics()).thenReturn(attemptMetric2);
    when(app2.getCurrentAppAttempt()).thenReturn(attempt2);

    rm.getRMContext().getRMApps().put(appId2, app2);

    SchedulerEvent addAppEvent2 =
        new AppAddedSchedulerEvent(appId2, "default", "user", null, Priority.newInstance(8));
    cs.handle(addAppEvent2);
    SchedulerEvent addAttemptEvent2 = new AppAttemptAddedSchedulerEvent(appAttemptId2, false);
    cs.handle(addAttemptEvent2);
    // add app end

    // Now, the first assignment will be for app2 since app2 is of highest
    // priority
    assertEquals(q.getApplications().size(), 2);
    assertEquals(q.getApplications().iterator().next().getApplicationAttemptId(), appAttemptId2);

    rm.stop();
  }
  @Test(timeout = 2000000)
  public void testUpdateHeartbeatResponseForCleanup() throws IOException {
    RMNodeImpl node = getRunningNode();
    NodeId nodeId = node.getNodeID();

    int rpcID = HopYarnAPIUtilities.getRPCID();
    byte[] allNMRequestData = new byte[1];
    allNMRequestData[0] = 0xA;
    try {
      RMUtilities.persistAppMasterRPC(rpcID, RPC.Type.RegisterNM, allNMRequestData);
    } catch (IOException ex) {
      Logger.getLogger(TestRMNodeTransitions.class.getName()).log(Level.SEVERE, null, ex);
    }
    TransactionState ts = new TransactionStateImpl(TransactionType.RM);

    // Expire a container
    ContainerId completedContainerId =
        BuilderUtils.newContainerId(
            BuilderUtils.newApplicationAttemptId(BuilderUtils.newApplicationId(0, 0), 0), 0);
    node.handle(new RMNodeCleanContainerEvent(nodeId, completedContainerId, ts));
    ts.decCounter(TransactionState.TransactionType.INIT);
    Assert.assertEquals(1, node.getContainersToCleanUp().size());

    rpcID = HopYarnAPIUtilities.getRPCID();
    allNMRequestData = new byte[1];
    allNMRequestData[0] = 0xA;
    try {
      RMUtilities.persistAppMasterRPC(rpcID, RPC.Type.FinishApplicationMaster, allNMRequestData);
    } catch (IOException ex) {
      Logger.getLogger(TestRMNodeTransitions.class.getName()).log(Level.SEVERE, null, ex);
    }
    TransactionState ts2 =
        new TransactionStateImpl(TransactionType.RM); // TransactionStateRM.newInstance(rpcID);

    // Finish an application
    ApplicationId finishedAppId = BuilderUtils.newApplicationId(0, 1);
    node.handle(new RMNodeCleanAppEvent(nodeId, finishedAppId, ts2));
    ts2.decCounter(TransactionState.TransactionType.INIT);
    Assert.assertEquals(1, node.getAppsToCleanup().size());
    rpcID = HopYarnAPIUtilities.getRPCID();
    allNMRequestData = new byte[1];
    allNMRequestData[0] = 0xA;
    try {
      RMUtilities.persistAppMasterRPC(rpcID, RPC.Type.FinishApplicationMaster, allNMRequestData);
    } catch (IOException ex) {
      Logger.getLogger(TestRMNodeTransitions.class.getName()).log(Level.SEVERE, null, ex);
    }
    TransactionState ts3 =
        new TransactionStateImpl(TransactionType.RM); // TransactionStateRM.newInstance(rpcID);

    // Verify status update does not clear containers/apps to cleanup
    // but updating heartbeat response for cleanup does
    RMNodeStatusEvent statusEvent = getMockRMNodeStatusEvent();
    RMNodeStatusEvent se =
        new RMNodeStatusEvent(
            statusEvent.getNodeId(),
            statusEvent.getNodeHealthStatus(),
            statusEvent.getContainers(),
            statusEvent.getKeepAliveAppIds(),
            statusEvent.getLatestResponse(),
            ts3);
    node.handle(se);

    ts3.decCounter(TransactionState.TransactionType.INIT);

    Assert.assertEquals(1, node.getContainersToCleanUp().size());
    Assert.assertEquals(1, node.getAppsToCleanup().size());
    NodeHeartbeatResponse hbrsp = Records.newRecord(NodeHeartbeatResponse.class);
    node.updateNodeHeartbeatResponseForCleanup(hbrsp, ts);
    Assert.assertEquals(0, node.getContainersToCleanUp().size());
    Assert.assertEquals(0, node.getAppsToCleanup().size());
    Assert.assertEquals(1, hbrsp.getContainersToCleanup().size());
    Assert.assertEquals(completedContainerId, hbrsp.getContainersToCleanup().get(0));
    Assert.assertEquals(1, hbrsp.getApplicationsToCleanup().size());
    Assert.assertEquals(finishedAppId, hbrsp.getApplicationsToCleanup().get(0));
  }
  @Test(timeout = 500000)
  public void testContainerUpdate() throws InterruptedException, IOException {
    // Start the node
    node.handle(
        new RMNodeEvent(
            null, RMNodeEventType.STARTED, new TransactionStateImpl(TransactionType.RM)));
    // If Distributed RT is enabled, this is the only way to let the scheduler
    // pick up the event, the PendingEvent retrieval does not invoke the
    // Mock Scheduler
    Configuration yarnconf = new YarnConfiguration();
    if (yarnconf.getBoolean(
        YarnConfiguration.DISTRIBUTED_RM, YarnConfiguration.DEFAULT_DISTRIBUTED_RM)) {
      scheduler.handle(
          new NodeAddedSchedulerEvent(node, new TransactionStateImpl(TransactionType.RM)));
    }
    NodeId nodeId = BuilderUtils.newNodeId("localhost:1", 1);
    RMNodeImpl node2 = new RMNodeImpl(nodeId, rmContext, "test", 0, 0, null, null, null);
    node2.handle(
        new RMNodeEvent(
            null, RMNodeEventType.STARTED, new TransactionStateImpl(TransactionType.RM)));
    // If Distributed RT is enabled, this is the only way to let the scheduler
    // pick up the event, the PendingEvent retrieval does not invoke the
    // Mock Scheduler
    if (yarnconf.getBoolean(
        YarnConfiguration.DISTRIBUTED_RM, YarnConfiguration.DEFAULT_DISTRIBUTED_RM)) {
      scheduler.handle(
          new NodeAddedSchedulerEvent(node2, new TransactionStateImpl(TransactionType.RM)));
    }
    ContainerId completedContainerIdFromNode1 =
        BuilderUtils.newContainerId(
            BuilderUtils.newApplicationAttemptId(BuilderUtils.newApplicationId(0, 0), 0), 0);
    ContainerId completedContainerIdFromNode2_1 =
        BuilderUtils.newContainerId(
            BuilderUtils.newApplicationAttemptId(BuilderUtils.newApplicationId(1, 1), 1), 1);
    ContainerId completedContainerIdFromNode2_2 =
        BuilderUtils.newContainerId(
            BuilderUtils.newApplicationAttemptId(BuilderUtils.newApplicationId(1, 1), 1), 2);

    RMNodeStatusEvent statusEventFromNode1 = getMockRMNodeStatusEvent();
    RMNodeStatusEvent statusEventFromNode2_1 = getMockRMNodeStatusEvent();
    RMNodeStatusEvent statusEventFromNode2_2 = getMockRMNodeStatusEvent();

    ContainerStatus containerStatusFromNode1 = mock(ContainerStatus.class);
    ContainerStatus containerStatusFromNode2_1 = mock(ContainerStatus.class);
    ContainerStatus containerStatusFromNode2_2 = mock(ContainerStatus.class);

    doReturn(completedContainerIdFromNode1).when(containerStatusFromNode1).getContainerId();
    doReturn(Collections.singletonList(containerStatusFromNode1))
        .when(statusEventFromNode1)
        .getContainers();
    doReturn(ContainerState.COMPLETE).when(containerStatusFromNode1).getState();
    node.handle(statusEventFromNode1);
    // If Distributed RT is enabled, this is the only way to let the scheduler
    // pick up the event, the PendingEvent retrieval does not invoke the
    // Mock Scheduler
    if (yarnconf.getBoolean(
        YarnConfiguration.DISTRIBUTED_RM, YarnConfiguration.DEFAULT_DISTRIBUTED_RM)) {
      scheduler.handle(
          new NodeUpdateSchedulerEvent(node, new TransactionStateImpl(TransactionType.RM)));
    }
    // ts.decCounter("test");
    Assert.assertEquals(1, completedContainers.size());
    Assert.assertEquals(completedContainerIdFromNode1, completedContainers.get(0).getContainerId());

    completedContainers.clear();

    doReturn(completedContainerIdFromNode2_1).when(containerStatusFromNode2_1).getContainerId();
    doReturn(Collections.singletonList(containerStatusFromNode2_1))
        .when(statusEventFromNode2_1)
        .getContainers();
    doReturn(ContainerState.COMPLETE).when(containerStatusFromNode2_1).getState();

    doReturn(completedContainerIdFromNode2_2).when(containerStatusFromNode2_2).getContainerId();
    doReturn(Collections.singletonList(containerStatusFromNode2_2))
        .when(statusEventFromNode2_2)
        .getContainers();
    doReturn(ContainerState.COMPLETE).when(containerStatusFromNode2_2).getState();

    node2.setNextHeartBeat(false);
    node2.handle(statusEventFromNode2_1);
    // If Distributed RT is enabled, this is the only way to let the scheduler
    // pick up the event, the PendingEvent retrieval does not invoke the
    // Mock Scheduler
    if (yarnconf.getBoolean(
        YarnConfiguration.DISTRIBUTED_RM, YarnConfiguration.DEFAULT_DISTRIBUTED_RM)) {
      scheduler.handle(
          new NodeUpdateSchedulerEvent(node2, new TransactionStateImpl(TransactionType.RM)));
    }

    node2.setNextHeartBeat(true);
    node2.handle(statusEventFromNode2_2);
    // If Distributed RT is enabled, this is the only way to let the scheduler
    // pick up the event, the PendingEvent retrieval does not invoke the
    // Mock Scheduler
    if (yarnconf.getBoolean(
        YarnConfiguration.DISTRIBUTED_RM, YarnConfiguration.DEFAULT_DISTRIBUTED_RM)) {
      scheduler.handle(
          new NodeUpdateSchedulerEvent(node2, new TransactionStateImpl(TransactionType.RM)));
    }
    // ts2.decCounter("test");
    Assert.assertEquals(2, completedContainers.size());
    Assert.assertEquals(
        completedContainerIdFromNode2_1, completedContainers.get(0).getContainerId());
    Assert.assertEquals(
        completedContainerIdFromNode2_2, completedContainers.get(1).getContainerId());
  }