@Test(timeout = 60000)
public void testNodeRemovalDuringAllocate() throws Exception {

  MockNM nm1 = new MockNM("h1:1234", 4096, rm.getResourceTrackerService());
  MockNM nm2 = new MockNM("h2:1234", 4096, rm.getResourceTrackerService());
  nm1.registerNode();
  nm2.registerNode();
  nm1.nodeHeartbeat(oppContainersStatus, true);
  nm2.nodeHeartbeat(oppContainersStatus, true);
  OpportunisticContainerAllocatorAMService amservice =
      (OpportunisticContainerAllocatorAMService) rm
          .getApplicationMasterService();
  MockRMAppSubmissionData data =
      MockRMAppSubmissionData.Builder.createWithMemory(1 * GB, rm)
          .withAppName("app")
          .withUser("user")
          .withAcls(null)
          .withQueue("default")
          .withUnmanagedAM(false)
          .build();
  RMApp app1 = MockRMAppSubmitter.submit(rm, data);
  ApplicationAttemptId attemptId =
      app1.getCurrentAppAttempt().getAppAttemptId();
  MockAM am1 = MockRM.launchAndRegisterAM(app1, rm, nm2);
  ResourceScheduler scheduler = rm.getResourceScheduler();
  RMNode rmNode1 = rm.getRMContext().getRMNodes().get(nm1.getNodeId());
  RMNode rmNode2 = rm.getRMContext().getRMNodes().get(nm2.getNodeId());
  OpportunisticContainerContext ctxt = ((CapacityScheduler) scheduler)
      .getApplicationAttempt(attemptId).getOpportunisticContainerContext();
  // Both node 1 and node 2 will be applicable for scheduling.
  nm1.nodeHeartbeat(oppContainersStatus, true);
  nm2.nodeHeartbeat(oppContainersStatus, true);
  for (int i = 0; i < 10; i++) {
    am1.allocate(
        Arrays.asList(ResourceRequest.newInstance(Priority.newInstance(1),
            "*", Resources.createResource(1 * GB), 2)),
        null);
    if (ctxt.getNodeMap().size() == 2) {
      break;
    }
    Thread.sleep(50);
  }
  Assert.assertEquals(2, ctxt.getNodeMap().size());
  // Remove node from scheduler but not from AM Service.
  scheduler.handle(new NodeRemovedSchedulerEvent(rmNode1));
  // After removal of node 1, only 1 node will be applicable for scheduling.
  for (int i = 0; i < 10; i++) {
    try {
      am1.allocate(
          Arrays.asList(ResourceRequest.newInstance(Priority.newInstance(1),
              "*", Resources.createResource(1 * GB), 2)),
          null);
    } catch (Exception e) {
      Assert.fail("Allocate request should be handled on node removal");
    }
    if (ctxt.getNodeMap().size() == 1) {
      break;
    }
    Thread.sleep(50);
  }
  Assert.assertEquals(1, ctxt.getNodeMap().size());
}