@Test
public void testNodeHeartbeatForAppCollectorsMap() throws Exception {

  Configuration conf = new Configuration();
  conf.setBoolean(YarnConfiguration.TIMELINE_SERVICE_ENABLED, true);
  // set version to 2
  conf.setFloat(YarnConfiguration.TIMELINE_SERVICE_VERSION, 2.0f);
  // enable aux-service based timeline collectors
  conf.set(YarnConfiguration.NM_AUX_SERVICES, "timeline_collector");
  conf.set(YarnConfiguration.NM_AUX_SERVICES + "."
      + "timeline_collector" + ".class",
      PerNodeTimelineCollectorsAuxService.class.getName());
  conf.setClass(YarnConfiguration.TIMELINE_SERVICE_WRITER_CLASS,
      FileSystemTimelineWriterImpl.class, TimelineWriter.class);
  rm = new MockRM(conf);
  rm.start();
  MockNM nm1 = rm.registerNode("host1:1234", 5120);
  MockNM nm2 = rm.registerNode("host2:1234", 2048);
  NodeHeartbeatResponse nodeHeartbeat1 = nm1.nodeHeartbeat(true);
  NodeHeartbeatResponse nodeHeartbeat2 = nm2.nodeHeartbeat(true);
  RMNodeImpl node1 =
      (RMNodeImpl) rm.getRMContext().getRMNodes().get(nm1.getNodeId());
  RMNodeImpl node2 =
      (RMNodeImpl) rm.getRMContext().getRMNodes().get(nm2.getNodeId());
  RMAppImpl app1 = (RMAppImpl) MockRMAppSubmitter.submitWithMemory(1024, rm);
  String collectorAddr1 = "1.2.3.4:5";
  app1.setCollectorData(AppCollectorData.newInstance(
      app1.getApplicationId(), collectorAddr1));
  String collectorAddr2 = "5.4.3.2:1";
  RMAppImpl app2 = (RMAppImpl) MockRMAppSubmitter.submitWithMemory(1024, rm);
  app2.setCollectorData(AppCollectorData.newInstance(
      app2.getApplicationId(), collectorAddr2));
  String collectorAddr3 = "5.4.3.2:2";
  app2.setCollectorData(AppCollectorData.newInstance(
      app2.getApplicationId(), collectorAddr3, 0, 1));
  String collectorAddr4 = "5.4.3.2:3";
  app2.setCollectorData(AppCollectorData.newInstance(
      app2.getApplicationId(), collectorAddr4, 1, 0));
  // Create a running container for app1 running on nm1
  ContainerId runningContainerId1 = BuilderUtils.newContainerId(
      BuilderUtils.newApplicationAttemptId(
      app1.getApplicationId(), 0), 0);
  ContainerStatus status1 = ContainerStatus.newInstance(runningContainerId1,
      ContainerState.RUNNING, "", 0);
  List<ContainerStatus> statusList = new ArrayList<ContainerStatus>();
  statusList.add(status1);
  NodeHealthStatus nodeHealth = NodeHealthStatus.newInstance(true,
      "", System.currentTimeMillis());
  NodeStatus nodeStatus = NodeStatus.newInstance(nm1.getNodeId(), 0,
      statusList, null, nodeHealth, null, null, null);
  node1.handle(new RMNodeStatusEvent(nm1.getNodeId(), nodeStatus));
  Assert.assertEquals(1, node1.getRunningApps().size());
  Assert.assertEquals(app1.getApplicationId(), node1.getRunningApps().get(0));
  // Create a running container for app2 running on nm2
  ContainerId runningContainerId2 = BuilderUtils.newContainerId(
      BuilderUtils.newApplicationAttemptId(
      app2.getApplicationId(), 0), 0);
  ContainerStatus status2 = ContainerStatus.newInstance(runningContainerId2,
      ContainerState.RUNNING, "", 0);
  statusList = new ArrayList<ContainerStatus>();
  statusList.add(status2);
  nodeStatus = NodeStatus.newInstance(nm1.getNodeId(), 0,
      statusList, null, nodeHealth, null, null, null);
  node2.handle(new RMNodeStatusEvent(nm2.getNodeId(), nodeStatus));
  Assert.assertEquals(1, node2.getRunningApps().size());
  Assert.assertEquals(app2.getApplicationId(), node2.getRunningApps().get(0));
  nodeHeartbeat1 = nm1.nodeHeartbeat(true);
  Map<ApplicationId, AppCollectorData> map1
      = nodeHeartbeat1.getAppCollectors();
  Assert.assertEquals(1, map1.size());
  Assert.assertEquals(collectorAddr1,
      map1.get(app1.getApplicationId()).getCollectorAddr());
  nodeHeartbeat2 = nm2.nodeHeartbeat(true);
  Map<ApplicationId, AppCollectorData> map2
      = nodeHeartbeat2.getAppCollectors();
  Assert.assertEquals(1, map2.size());
  Assert.assertEquals(collectorAddr4,
      map2.get(app2.getApplicationId()).getCollectorAddr());
}