@Test
public void testGetApplicableNodeCountForAMLabels() throws Exception {

  Set<NodeId> noLabelNodes = new HashSet<>();
  for (int i = 0; i < 80; i++) {
    noLabelNodes.add(NodeId.newInstance("host" + i, 1234));
  }
  Set<NodeId> label1Nodes = new HashSet<>();
  for (int i = 80; i < 90; i++) {
    label1Nodes.add(NodeId.newInstance("host" + i, 1234));
  }
  label1Nodes.add(NodeId.newInstance("host101", 0));
  label1Nodes.add(NodeId.newInstance("host102", 0));
  Map<String, Set<NodeId>> label1NodesMap = new HashMap<>();
  label1NodesMap.put("label1", label1Nodes);
  YarnConfiguration conf = new YarnConfiguration();
  conf.setBoolean(YarnConfiguration.NODE_LABELS_ENABLED, true);
  ResourceScheduler scheduler = Mockito.mock(ResourceScheduler.class);
  Mockito.when(scheduler.getNumClusterNodes()).thenReturn(100);
  RMContext rmContext = Mockito.mock(RMContext.class);
  Mockito.when(rmContext.getScheduler()).thenReturn(scheduler);
  RMNodeLabelsManager labMan = Mockito.mock(RMNodeLabelsManager.class);
  Mockito.when(labMan.getNodesWithoutALabel()).thenReturn(noLabelNodes);
  Mockito.when(labMan.getLabelsToNodes(Collections.singleton("label1")))
      .thenReturn(label1NodesMap);
  Mockito.when(rmContext.getNodeLabelManager()).thenReturn(labMan);
  ResourceRequest anyReq = createResourceRequest(ResourceRequest.ANY,
      true, null);
  List<ResourceRequest> reqs = new ArrayList<>();
  reqs.add(anyReq);
  Assert.assertEquals(80,
      RMServerUtils.getApplicableNodeCountForAM(rmContext, conf, reqs));
  anyReq.setNodeLabelExpression("label1");
  Assert.assertEquals(10,
      RMServerUtils.getApplicableNodeCountForAM(rmContext, conf, reqs));
}