@Test
public void testGetApplicableNodeCountForAMLocalityAndLabels()
    throws Exception {

  List<NodeId> rack1Nodes = new ArrayList<>();
  for (int i = 0; i < 29; i++) {
    rack1Nodes.add(NodeId.newInstance("host" + i, 1234));
  }
  NodeId node1 = NodeId.newInstance("node1", 1234);
  NodeId node2 = NodeId.newInstance("node2", 1234);
  rack1Nodes.add(node2);
  Set<NodeId> noLabelNodes = new HashSet<>();
  for (int i = 0; i < 19; i++) {
    noLabelNodes.add(rack1Nodes.get(i));
  }
  noLabelNodes.add(node2);
  for (int i = 29; i < 89; i++) {
    noLabelNodes.add(NodeId.newInstance("host" + i, 1234));
  }
  Set<NodeId> label1Nodes = new HashSet<>();
  label1Nodes.add(node1);
  for (int i = 89; i < 93; i++) {
    label1Nodes.add(NodeId.newInstance("host" + i, 1234));
  }
  for (int i = 19; i < 29; i++) {
    label1Nodes.add(rack1Nodes.get(i));
  }
  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);
  Mockito.when(scheduler.getNodeIds("/rack1")).thenReturn(rack1Nodes);
  Mockito.when(scheduler.getNodeIds("node1"))
      .thenReturn(Collections.singletonList(node1));
  Mockito.when(scheduler.getNodeIds("node2"))
      .thenReturn(Collections.singletonList(node2));
  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));
  ResourceRequest rackReq = createResourceRequest("/rack1", true, null);
  reqs.add(rackReq);
  Assert.assertEquals(20,
      RMServerUtils.getApplicableNodeCountForAM(rmContext, conf, reqs));
  anyReq.setRelaxLocality(false);
  Assert.assertEquals(20,
      RMServerUtils.getApplicableNodeCountForAM(rmContext, conf, reqs));
  rackReq.setRelaxLocality(false);
  Assert.assertEquals(80,
      RMServerUtils.getApplicableNodeCountForAM(rmContext, conf, reqs));
  ResourceRequest node1Req = createResourceRequest("node1", false, null);
  reqs.add(node1Req);
  Assert.assertEquals(80,
      RMServerUtils.getApplicableNodeCountForAM(rmContext, conf, reqs));
  node1Req.setRelaxLocality(true);
  Assert.assertEquals(0,
      RMServerUtils.getApplicableNodeCountForAM(rmContext, conf, reqs));
  rackReq.setRelaxLocality(true);
  Assert.assertEquals(20,
      RMServerUtils.getApplicableNodeCountForAM(rmContext, conf, reqs));
  ResourceRequest node2Req = createResourceRequest("node2", false, null);
  reqs.add(node2Req);
  Assert.assertEquals(20,
      RMServerUtils.getApplicableNodeCountForAM(rmContext, conf, reqs));
  node2Req.setRelaxLocality(true);
  Assert.assertEquals(20,
      RMServerUtils.getApplicableNodeCountForAM(rmContext, conf, reqs));
  rackReq.setRelaxLocality(false);
  Assert.assertEquals(1,
      RMServerUtils.getApplicableNodeCountForAM(rmContext, conf, reqs));
  node1Req.setRelaxLocality(false);
  Assert.assertEquals(1,
      RMServerUtils.getApplicableNodeCountForAM(rmContext, conf, reqs));
  node2Req.setRelaxLocality(false);
  Assert.assertEquals(80,
      RMServerUtils.getApplicableNodeCountForAM(rmContext, conf, reqs));
  anyReq.setNodeLabelExpression("label1");
  rackReq.setNodeLabelExpression("label1");
  node1Req.setNodeLabelExpression("label1");
  node2Req.setNodeLabelExpression("label1");
  anyReq.setRelaxLocality(true);
  reqs = new ArrayList<>();
  reqs.add(anyReq);
  Assert.assertEquals(15,
      RMServerUtils.getApplicableNodeCountForAM(rmContext, conf, reqs));
  rackReq.setRelaxLocality(true);
  reqs.add(rackReq);
  Assert.assertEquals(10,
      RMServerUtils.getApplicableNodeCountForAM(rmContext, conf, reqs));
  anyReq.setRelaxLocality(false);
  Assert.assertEquals(10,
      RMServerUtils.getApplicableNodeCountForAM(rmContext, conf, reqs));
  rackReq.setRelaxLocality(false);
  Assert.assertEquals(15,
      RMServerUtils.getApplicableNodeCountForAM(rmContext, conf, reqs));
  node1Req.setRelaxLocality(false);
  reqs.add(node1Req);
  Assert.assertEquals(15,
      RMServerUtils.getApplicableNodeCountForAM(rmContext, conf, reqs));
  node1Req.setRelaxLocality(true);
  Assert.assertEquals(1,
      RMServerUtils.getApplicableNodeCountForAM(rmContext, conf, reqs));
  rackReq.setRelaxLocality(true);
  Assert.assertEquals(11,
      RMServerUtils.getApplicableNodeCountForAM(rmContext, conf, reqs));
  node2Req.setRelaxLocality(false);
  reqs.add(node2Req);
  Assert.assertEquals(11,
      RMServerUtils.getApplicableNodeCountForAM(rmContext, conf, reqs));
  node2Req.setRelaxLocality(true);
  Assert.assertEquals(11,
      RMServerUtils.getApplicableNodeCountForAM(rmContext, conf, reqs));
  rackReq.setRelaxLocality(false);
  Assert.assertEquals(1,
      RMServerUtils.getApplicableNodeCountForAM(rmContext, conf, reqs));
  node1Req.setRelaxLocality(false);
  Assert.assertEquals(0,
      RMServerUtils.getApplicableNodeCountForAM(rmContext, conf, reqs));
  node2Req.setRelaxLocality(false);
  Assert.assertEquals(15,
      RMServerUtils.getApplicableNodeCountForAM(rmContext, conf, reqs));
}