org.apache.hadoop.yarn.server.resourcemanager.metrics.SystemMetricsPublisher Java Examples

@Override
@Before
public void setUp() throws Exception {
  super.setUp();
  context = new MockRMContext();
  context.setDispatcher(TestObjectFactory.getMockDispatcher());
  context.setSystemMetricsPublisher(new SystemMetricsPublisher());
  
  nodeOne = TestObjectFactory.getRMNode("localhost-one", 8800, Resource.newInstance(1024, 2));
  nodeTwo = TestObjectFactory.getRMNode("localhost-two", 8800, Resource.newInstance(2048, 4));
  sNodeOne = new FSSchedulerNode(nodeOne, false);
  sNodeTwo = new FSSchedulerNode(nodeTwo, false);

  store = new NodeStore();
  store.add(sNodeOne);
  store.add(sNodeTwo);

  containerOne = TestObjectFactory.getRMContainer(nodeOne, context, 1, 2, 1024);
}
 

private void mockRMContext(YarnScheduler yarnScheduler, RMContext rmContext)
    throws IOException {
  Dispatcher dispatcher = mock(Dispatcher.class);
  when(rmContext.getDispatcher()).thenReturn(dispatcher);
  EventHandler eventHandler = mock(EventHandler.class);
  when(dispatcher.getEventHandler()).thenReturn(eventHandler);
  QueueInfo queInfo = recordFactory.newRecordInstance(QueueInfo.class);
  queInfo.setQueueName("testqueue");
  when(yarnScheduler.getQueueInfo(eq("testqueue"), anyBoolean(), anyBoolean()))
      .thenReturn(queInfo);
  when(yarnScheduler.getQueueInfo(eq("nonexistentqueue"), anyBoolean(), anyBoolean()))
      .thenThrow(new IOException("queue does not exist"));
  RMApplicationHistoryWriter writer = mock(RMApplicationHistoryWriter.class);
  when(rmContext.getRMApplicationHistoryWriter()).thenReturn(writer);
  SystemMetricsPublisher publisher = mock(SystemMetricsPublisher.class);
  when(rmContext.getSystemMetricsPublisher()).thenReturn(publisher);
  ConcurrentHashMap<ApplicationId, RMApp> apps = getRMApps(rmContext,
      yarnScheduler);
  when(rmContext.getRMApps()).thenReturn(apps);
  when(yarnScheduler.getAppsInQueue(eq("testqueue"))).thenReturn(
      getSchedulerApps(apps));
   ResourceScheduler rs = mock(ResourceScheduler.class);
   when(rmContext.getScheduler()).thenReturn(rs);
}
 

private void mockRMContext(YarnScheduler yarnScheduler, RMContext rmContext)
    throws IOException {
  Dispatcher dispatcher = mock(Dispatcher.class);
  when(rmContext.getDispatcher()).thenReturn(dispatcher);
  EventHandler eventHandler = mock(EventHandler.class);
  when(dispatcher.getEventHandler()).thenReturn(eventHandler);
  QueueInfo queInfo = recordFactory.newRecordInstance(QueueInfo.class);
  queInfo.setQueueName("testqueue");
  when(yarnScheduler.getQueueInfo(eq("testqueue"), anyBoolean(), anyBoolean()))
      .thenReturn(queInfo);
  when(yarnScheduler.getQueueInfo(eq("nonexistentqueue"), anyBoolean(), anyBoolean()))
      .thenThrow(new IOException("queue does not exist"));
  RMApplicationHistoryWriter writer = mock(RMApplicationHistoryWriter.class);
  when(rmContext.getRMApplicationHistoryWriter()).thenReturn(writer);
  SystemMetricsPublisher publisher = mock(SystemMetricsPublisher.class);
  when(rmContext.getSystemMetricsPublisher()).thenReturn(publisher);
  ConcurrentHashMap<ApplicationId, RMApp> apps = getRMApps(rmContext,
      yarnScheduler);
  when(rmContext.getRMApps()).thenReturn(apps);
  when(yarnScheduler.getAppsInQueue(eq("testqueue"))).thenReturn(
      getSchedulerApps(apps));
   ResourceScheduler rs = mock(ResourceScheduler.class);
   when(rmContext.getScheduler()).thenReturn(rs);
}
 

@Override
@Before
public void setUp() throws Exception {
  super.setUp();
  this.baseStateStoreDirectory = "/tmp/yarn-node-capacity-manager-test";
  context = new MockRMContext();
  context.setDispatcher(TestObjectFactory.getMockDispatcher());
  context.setSystemMetricsPublisher(new SystemMetricsPublisher());
  
  nodeOne = TestObjectFactory.getRMNode("localhost-one", 8800, Resource.newInstance(2048, 4));
  nodeTwo = TestObjectFactory.getRMNode("localhost-two", 8800, Resource.newInstance(1024, 2));
  sNodeOne = new FSSchedulerNode(nodeOne, false);
  sNodeTwo = new FSSchedulerNode(nodeTwo, false);

  containerOne = TestObjectFactory.getRMContainer(nodeOne, context, 1, 2, 1024);
  store = new NodeStore();
  store.add(sNodeOne);
  store.add(sNodeTwo);
  
  MyriadDriver driver = TestObjectFactory.getMyriadDriver(new MockSchedulerDriver());
  olManager = new OfferLifecycleManager(store, driver);
  state = TestObjectFactory.getSchedulerState(new MyriadConfiguration(), "/tmp/yarn-node-capacity-manager-test");
  MyriadFairScheduler scheduler = TestObjectFactory.getMyriadFairScheduler(context);
  
  scheduler.addNode(sNodeOne);
  scheduler.addNode(sNodeTwo);
  manager = new YarnNodeCapacityManager(new CompositeInterceptor(), scheduler, 
            context, driver, olManager, store, state, new TaskUtils(this.cfg));
}
 

public RMContext mockRMContext(int n, long time) {
  final List<RMApp> apps = newRMApps(n, time, RMAppState.FINISHED);
  final ConcurrentMap<ApplicationId, RMApp> map = Maps.newConcurrentMap();
  for (RMApp app : apps) {
    map.put(app.getApplicationId(), app);
  }
  Dispatcher rmDispatcher = new AsyncDispatcher();
  ContainerAllocationExpirer containerAllocationExpirer = new ContainerAllocationExpirer(
      rmDispatcher);
  AMLivelinessMonitor amLivelinessMonitor = new AMLivelinessMonitor(
      rmDispatcher);
  AMLivelinessMonitor amFinishingMonitor = new AMLivelinessMonitor(
      rmDispatcher);
  RMApplicationHistoryWriter writer = mock(RMApplicationHistoryWriter.class);
  RMContext context = new RMContextImpl(rmDispatcher,
      containerAllocationExpirer, amLivelinessMonitor, amFinishingMonitor,
      null, null, null, null, null, writer) {
    @Override
    public ConcurrentMap<ApplicationId, RMApp> getRMApps() {
      return map;
    }
  };
  ((RMContextImpl)context).setStateStore(mock(RMStateStore.class));
  metricsPublisher = mock(SystemMetricsPublisher.class);
  ((RMContextImpl)context).setSystemMetricsPublisher(metricsPublisher);
  return context;
}
 

public RMContext mockRMContext(int n, long time) {
  final List<RMApp> apps = newRMApps(n, time, RMAppState.FINISHED);
  final ConcurrentMap<ApplicationId, RMApp> map = Maps.newConcurrentMap();
  for (RMApp app : apps) {
    map.put(app.getApplicationId(), app);
  }
  Dispatcher rmDispatcher = new AsyncDispatcher();
  ContainerAllocationExpirer containerAllocationExpirer = new ContainerAllocationExpirer(
      rmDispatcher);
  AMLivelinessMonitor amLivelinessMonitor = new AMLivelinessMonitor(
      rmDispatcher);
  AMLivelinessMonitor amFinishingMonitor = new AMLivelinessMonitor(
      rmDispatcher);
  RMApplicationHistoryWriter writer = mock(RMApplicationHistoryWriter.class);
  RMContext context = new RMContextImpl(rmDispatcher,
      containerAllocationExpirer, amLivelinessMonitor, amFinishingMonitor,
      null, null, null, null, null, writer) {
    @Override
    public ConcurrentMap<ApplicationId, RMApp> getRMApps() {
      return map;
    }
  };
  ((RMContextImpl)context).setStateStore(mock(RMStateStore.class));
  metricsPublisher = mock(SystemMetricsPublisher.class);
  ((RMContextImpl)context).setSystemMetricsPublisher(metricsPublisher);
  return context;
}
 

@Test(timeout=5000)
public void testAppAttemptMetrics() throws Exception {
  AsyncDispatcher dispatcher = new InlineDispatcher();
  
  FifoScheduler scheduler = new FifoScheduler();
  RMApplicationHistoryWriter writer = mock(RMApplicationHistoryWriter.class);
  RMContext rmContext = new RMContextImpl(dispatcher, null,
      null, null, null, null, null, null, null, writer, scheduler);
  ((RMContextImpl) rmContext).setSystemMetricsPublisher(
      mock(SystemMetricsPublisher.class));

  Configuration conf = new Configuration();
  scheduler.setRMContext(rmContext);
  scheduler.init(conf);
  scheduler.start();
  scheduler.reinitialize(conf, rmContext);
  QueueMetrics metrics = scheduler.getRootQueueMetrics();
  int beforeAppsSubmitted = metrics.getAppsSubmitted();

  ApplicationId appId = BuilderUtils.newApplicationId(200, 1);
  ApplicationAttemptId appAttemptId = BuilderUtils.newApplicationAttemptId(
      appId, 1);

  SchedulerEvent appEvent = new AppAddedSchedulerEvent(appId, "queue", "user");
  scheduler.handle(appEvent);
  SchedulerEvent attemptEvent =
      new AppAttemptAddedSchedulerEvent(appAttemptId, false);
  scheduler.handle(attemptEvent);

  appAttemptId = BuilderUtils.newApplicationAttemptId(appId, 2);
  SchedulerEvent attemptEvent2 =
      new AppAttemptAddedSchedulerEvent(appAttemptId, false);
  scheduler.handle(attemptEvent2);

  int afterAppsSubmitted = metrics.getAppsSubmitted();
  Assert.assertEquals(1, afterAppsSubmitted - beforeAppsSubmitted);
  scheduler.stop();
}
 

@Test(timeout=5000)
public void testAppAttemptMetrics() throws Exception {
  AsyncDispatcher dispatcher = new InlineDispatcher();
  
  FifoScheduler scheduler = new FifoScheduler();
  RMApplicationHistoryWriter writer = mock(RMApplicationHistoryWriter.class);
  RMContext rmContext = new RMContextImpl(dispatcher, null,
      null, null, null, null, null, null, null, writer, scheduler);
  ((RMContextImpl) rmContext).setSystemMetricsPublisher(
      mock(SystemMetricsPublisher.class));

  Configuration conf = new Configuration();
  scheduler.setRMContext(rmContext);
  scheduler.init(conf);
  scheduler.start();
  scheduler.reinitialize(conf, rmContext);
  QueueMetrics metrics = scheduler.getRootQueueMetrics();
  int beforeAppsSubmitted = metrics.getAppsSubmitted();

  ApplicationId appId = BuilderUtils.newApplicationId(200, 1);
  ApplicationAttemptId appAttemptId = BuilderUtils.newApplicationAttemptId(
      appId, 1);

  SchedulerEvent appEvent = new AppAddedSchedulerEvent(appId, "queue", "user");
  scheduler.handle(appEvent);
  SchedulerEvent attemptEvent =
      new AppAttemptAddedSchedulerEvent(appAttemptId, false);
  scheduler.handle(attemptEvent);

  appAttemptId = BuilderUtils.newApplicationAttemptId(appId, 2);
  SchedulerEvent attemptEvent2 =
      new AppAttemptAddedSchedulerEvent(appAttemptId, false);
  scheduler.handle(attemptEvent2);

  int afterAppsSubmitted = metrics.getAppsSubmitted();
  Assert.assertEquals(1, afterAppsSubmitted - beforeAppsSubmitted);
  scheduler.stop();
}
 

@Test(timeout=2000)
public void testNodeLocalAssignment() throws Exception {
  AsyncDispatcher dispatcher = new InlineDispatcher();
  Configuration conf = new Configuration();
  RMContainerTokenSecretManager containerTokenSecretManager =
      new RMContainerTokenSecretManager(conf);
  containerTokenSecretManager.rollMasterKey();
  NMTokenSecretManagerInRM nmTokenSecretManager =
      new NMTokenSecretManagerInRM(conf);
  nmTokenSecretManager.rollMasterKey();
  RMApplicationHistoryWriter writer = mock(RMApplicationHistoryWriter.class);
  
  FifoScheduler scheduler = new FifoScheduler();
  RMContext rmContext = new RMContextImpl(dispatcher, null, null, null, null,
      null, containerTokenSecretManager, nmTokenSecretManager, null, writer,
      scheduler);
  ((RMContextImpl) rmContext).setSystemMetricsPublisher(
      mock(SystemMetricsPublisher.class));

  scheduler.setRMContext(rmContext);
  scheduler.init(conf);
  scheduler.start();
  scheduler.reinitialize(new Configuration(), rmContext);

  RMNode node0 = MockNodes.newNodeInfo(1,
      Resources.createResource(1024 * 64), 1, "127.0.0.1");
  NodeAddedSchedulerEvent nodeEvent1 = new NodeAddedSchedulerEvent(node0);
  scheduler.handle(nodeEvent1);

  int _appId = 1;
  int _appAttemptId = 1;
  ApplicationAttemptId appAttemptId = createAppAttemptId(_appId,
      _appAttemptId);

  createMockRMApp(appAttemptId, rmContext);

  AppAddedSchedulerEvent appEvent =
      new AppAddedSchedulerEvent(appAttemptId.getApplicationId(), "queue1",
          "user1");
  scheduler.handle(appEvent);
  AppAttemptAddedSchedulerEvent attemptEvent =
      new AppAttemptAddedSchedulerEvent(appAttemptId, false);
  scheduler.handle(attemptEvent);

  int memory = 64;
  int nConts = 3;
  int priority = 20;

  List<ResourceRequest> ask = new ArrayList<ResourceRequest>();
  ResourceRequest nodeLocal = createResourceRequest(memory,
      node0.getHostName(), priority, nConts);
  ResourceRequest rackLocal = createResourceRequest(memory,
      node0.getRackName(), priority, nConts);
  ResourceRequest any = createResourceRequest(memory, ResourceRequest.ANY, priority,
      nConts);
  ask.add(nodeLocal);
  ask.add(rackLocal);
  ask.add(any);
  scheduler.allocate(appAttemptId, ask, new ArrayList<ContainerId>(), null, null);

  NodeUpdateSchedulerEvent node0Update = new NodeUpdateSchedulerEvent(node0);

  // Before the node update event, there are 3 local requests outstanding
  Assert.assertEquals(3, nodeLocal.getNumContainers());

  scheduler.handle(node0Update);

  // After the node update event, check that there are no more local requests
  // outstanding
  Assert.assertEquals(0, nodeLocal.getNumContainers());
  //Also check that the containers were scheduled
  SchedulerAppReport info = scheduler.getSchedulerAppInfo(appAttemptId);
  Assert.assertEquals(3, info.getLiveContainers().size());
  scheduler.stop();
}
 

@Override
public void setSystemMetricsPublisher(
    SystemMetricsPublisher systemMetricsPublisher) {
  activeServiceContext.setSystemMetricsPublisher(systemMetricsPublisher);
}
 

@Override
public SystemMetricsPublisher getSystemMetricsPublisher() {
  return activeServiceContext.getSystemMetricsPublisher();
}
 

@Test
public void testReleaseWhileRunning() {

  DrainDispatcher drainDispatcher = new DrainDispatcher();
  EventHandler<RMAppAttemptEvent> appAttemptEventHandler = mock(EventHandler.class);
  EventHandler generic = mock(EventHandler.class);
  drainDispatcher.register(RMAppAttemptEventType.class,
      appAttemptEventHandler);
  drainDispatcher.register(RMNodeEventType.class, generic);
  drainDispatcher.init(new YarnConfiguration());
  drainDispatcher.start();
  NodeId nodeId = BuilderUtils.newNodeId("host", 3425);
  ApplicationId appId = BuilderUtils.newApplicationId(1, 1);
  ApplicationAttemptId appAttemptId = BuilderUtils.newApplicationAttemptId(
      appId, 1);
  ContainerId containerId = BuilderUtils.newContainerId(appAttemptId, 1);
  ContainerAllocationExpirer expirer = mock(ContainerAllocationExpirer.class);

  Resource resource = BuilderUtils.newResource(512, 1);
  Priority priority = BuilderUtils.newPriority(5);

  Container container = BuilderUtils.newContainer(containerId, nodeId,
      "host:3465", resource, priority, null);
  ConcurrentMap<ApplicationId, RMApp> rmApps =
      spy(new ConcurrentHashMap<ApplicationId, RMApp>());
  RMApp rmApp = mock(RMApp.class);
  when(rmApp.getRMAppAttempt((ApplicationAttemptId)Matchers.any())).thenReturn(null);
  Mockito.doReturn(rmApp).when(rmApps).get((ApplicationId)Matchers.any());

  RMApplicationHistoryWriter writer = mock(RMApplicationHistoryWriter.class);
  SystemMetricsPublisher publisher = mock(SystemMetricsPublisher.class);
  RMContext rmContext = mock(RMContext.class);
  when(rmContext.getDispatcher()).thenReturn(drainDispatcher);
  when(rmContext.getContainerAllocationExpirer()).thenReturn(expirer);
  when(rmContext.getRMApplicationHistoryWriter()).thenReturn(writer);
  when(rmContext.getRMApps()).thenReturn(rmApps);
  when(rmContext.getSystemMetricsPublisher()).thenReturn(publisher);
  when(rmContext.getYarnConfiguration()).thenReturn(new YarnConfiguration());
  RMContainer rmContainer = new RMContainerImpl(container, appAttemptId,
      nodeId, "user", rmContext);

  assertEquals(RMContainerState.NEW, rmContainer.getState());
  assertEquals(resource, rmContainer.getAllocatedResource());
  assertEquals(nodeId, rmContainer.getAllocatedNode());
  assertEquals(priority, rmContainer.getAllocatedPriority());
  verify(writer).containerStarted(any(RMContainer.class));
  verify(publisher).containerCreated(any(RMContainer.class), anyLong());

  rmContainer.handle(new RMContainerEvent(containerId,
      RMContainerEventType.START));
  drainDispatcher.await();
  assertEquals(RMContainerState.ALLOCATED, rmContainer.getState());
  rmContainer.handle(new RMContainerEvent(containerId,
      RMContainerEventType.ACQUIRED));
  drainDispatcher.await();
  assertEquals(RMContainerState.ACQUIRED, rmContainer.getState());

  rmContainer.handle(new RMContainerEvent(containerId,
      RMContainerEventType.LAUNCHED));
  drainDispatcher.await();
  assertEquals(RMContainerState.RUNNING, rmContainer.getState());
  assertEquals("http://host:3465/node/containerlogs/container_1_0001_01_000001/user",
      rmContainer.getLogURL());

  // In RUNNING state. Verify RELEASED and associated actions.
  reset(appAttemptEventHandler);
  ContainerStatus containerStatus = SchedulerUtils
      .createAbnormalContainerStatus(containerId,
          SchedulerUtils.RELEASED_CONTAINER);
  rmContainer.handle(new RMContainerFinishedEvent(containerId,
      containerStatus, RMContainerEventType.RELEASED));
  drainDispatcher.await();
  assertEquals(RMContainerState.RELEASED, rmContainer.getState());
  assertEquals(SchedulerUtils.RELEASED_CONTAINER,
      rmContainer.getDiagnosticsInfo());
  assertEquals(ContainerExitStatus.ABORTED,
      rmContainer.getContainerExitStatus());
  assertEquals(ContainerState.COMPLETE, rmContainer.getContainerState());
  verify(writer).containerFinished(any(RMContainer.class));
  verify(publisher).containerFinished(any(RMContainer.class), anyLong());

  ArgumentCaptor<RMAppAttemptContainerFinishedEvent> captor = ArgumentCaptor
      .forClass(RMAppAttemptContainerFinishedEvent.class);
  verify(appAttemptEventHandler).handle(captor.capture());
  RMAppAttemptContainerFinishedEvent cfEvent = captor.getValue();
  assertEquals(appAttemptId, cfEvent.getApplicationAttemptId());
  assertEquals(containerStatus, cfEvent.getContainerStatus());
  assertEquals(RMAppAttemptEventType.CONTAINER_FINISHED, cfEvent.getType());
  
  // In RELEASED state. A FINIHSED event may come in.
  rmContainer.handle(new RMContainerFinishedEvent(containerId, SchedulerUtils
      .createAbnormalContainerStatus(containerId, "FinishedContainer"),
      RMContainerEventType.FINISHED));
  assertEquals(RMContainerState.RELEASED, rmContainer.getState());
}
 

@Test
public void testExpireWhileRunning() {

  DrainDispatcher drainDispatcher = new DrainDispatcher();
  EventHandler<RMAppAttemptEvent> appAttemptEventHandler = mock(EventHandler.class);
  EventHandler generic = mock(EventHandler.class);
  drainDispatcher.register(RMAppAttemptEventType.class,
      appAttemptEventHandler);
  drainDispatcher.register(RMNodeEventType.class, generic);
  drainDispatcher.init(new YarnConfiguration());
  drainDispatcher.start();
  NodeId nodeId = BuilderUtils.newNodeId("host", 3425);
  ApplicationId appId = BuilderUtils.newApplicationId(1, 1);
  ApplicationAttemptId appAttemptId = BuilderUtils.newApplicationAttemptId(
      appId, 1);
  ContainerId containerId = BuilderUtils.newContainerId(appAttemptId, 1);
  ContainerAllocationExpirer expirer = mock(ContainerAllocationExpirer.class);

  Resource resource = BuilderUtils.newResource(512, 1);
  Priority priority = BuilderUtils.newPriority(5);

  Container container = BuilderUtils.newContainer(containerId, nodeId,
      "host:3465", resource, priority, null);

  RMApplicationHistoryWriter writer = mock(RMApplicationHistoryWriter.class);
  SystemMetricsPublisher publisher = mock(SystemMetricsPublisher.class);
  RMContext rmContext = mock(RMContext.class);
  when(rmContext.getDispatcher()).thenReturn(drainDispatcher);
  when(rmContext.getContainerAllocationExpirer()).thenReturn(expirer);
  when(rmContext.getRMApplicationHistoryWriter()).thenReturn(writer);
  when(rmContext.getSystemMetricsPublisher()).thenReturn(publisher);
  when(rmContext.getYarnConfiguration()).thenReturn(new YarnConfiguration());
  RMContainer rmContainer = new RMContainerImpl(container, appAttemptId,
      nodeId, "user", rmContext);

  assertEquals(RMContainerState.NEW, rmContainer.getState());
  assertEquals(resource, rmContainer.getAllocatedResource());
  assertEquals(nodeId, rmContainer.getAllocatedNode());
  assertEquals(priority, rmContainer.getAllocatedPriority());
  verify(writer).containerStarted(any(RMContainer.class));
  verify(publisher).containerCreated(any(RMContainer.class), anyLong());

  rmContainer.handle(new RMContainerEvent(containerId,
      RMContainerEventType.START));
  drainDispatcher.await();
  assertEquals(RMContainerState.ALLOCATED, rmContainer.getState());

  rmContainer.handle(new RMContainerEvent(containerId,
      RMContainerEventType.ACQUIRED));
  drainDispatcher.await();
  assertEquals(RMContainerState.ACQUIRED, rmContainer.getState());

  rmContainer.handle(new RMContainerEvent(containerId,
      RMContainerEventType.LAUNCHED));
  drainDispatcher.await();
  assertEquals(RMContainerState.RUNNING, rmContainer.getState());
  assertEquals("http://host:3465/node/containerlogs/container_1_0001_01_000001/user",
      rmContainer.getLogURL());

  // In RUNNING state. Verify EXPIRE and associated actions.
  reset(appAttemptEventHandler);
  ContainerStatus containerStatus = SchedulerUtils
      .createAbnormalContainerStatus(containerId,
          SchedulerUtils.EXPIRED_CONTAINER);
  rmContainer.handle(new RMContainerFinishedEvent(containerId,
      containerStatus, RMContainerEventType.EXPIRE));
  drainDispatcher.await();
  assertEquals(RMContainerState.RUNNING, rmContainer.getState());
  verify(writer, never()).containerFinished(any(RMContainer.class));
  verify(publisher, never()).containerFinished(any(RMContainer.class),
      anyLong());
}
 

@Before
public void setUp() throws Exception {
  conf = new YarnConfiguration();
  AuthenticationMethod authMethod = AuthenticationMethod.SIMPLE;
  if (isSecurityEnabled) {
    authMethod = AuthenticationMethod.KERBEROS;
  }
  SecurityUtil.setAuthenticationMethod(authMethod, conf);
  UserGroupInformation.setConfiguration(conf);

  rmDispatcher = new DrainDispatcher();
  ContainerAllocationExpirer containerAllocationExpirer = 
      mock(ContainerAllocationExpirer.class);
  AMLivelinessMonitor amLivelinessMonitor = mock(AMLivelinessMonitor.class);
  AMLivelinessMonitor amFinishingMonitor = mock(AMLivelinessMonitor.class);
  store = mock(RMStateStore.class);
  writer = mock(RMApplicationHistoryWriter.class);
  DelegationTokenRenewer renewer = mock(DelegationTokenRenewer.class);
  RMContext realRMContext = 
      new RMContextImpl(rmDispatcher,
        containerAllocationExpirer, amLivelinessMonitor, amFinishingMonitor,
        renewer, new AMRMTokenSecretManager(conf, this.rmContext),
        new RMContainerTokenSecretManager(conf),
        new NMTokenSecretManagerInRM(conf),
        new ClientToAMTokenSecretManagerInRM(),
        writer);
  ((RMContextImpl)realRMContext).setStateStore(store);
  publisher = mock(SystemMetricsPublisher.class);
  ((RMContextImpl)realRMContext).setSystemMetricsPublisher(publisher);

  this.rmContext = spy(realRMContext);

  ResourceScheduler resourceScheduler = mock(ResourceScheduler.class);
  doReturn(null).when(resourceScheduler)
            .getAppResourceUsageReport((ApplicationAttemptId)Matchers.any());
  doReturn(resourceScheduler).when(rmContext).getScheduler();

  rmDispatcher.register(RMAppAttemptEventType.class,
      new TestApplicationAttemptEventDispatcher(this.rmContext));

  rmDispatcher.register(RMAppEventType.class,
      new TestApplicationEventDispatcher(rmContext));
  
  rmDispatcher.register(RMAppManagerEventType.class,
      new TestApplicationManagerEventDispatcher());
  
  schedulerDispatcher = new TestSchedulerEventDispatcher();
  rmDispatcher.register(SchedulerEventType.class,
      schedulerDispatcher);
  
  rmDispatcher.init(conf);
  rmDispatcher.start();
}
 

@SuppressWarnings("deprecation")
@Before
public void setUp() throws Exception {
  AuthenticationMethod authMethod = AuthenticationMethod.SIMPLE;
  if (isSecurityEnabled) {
    authMethod = AuthenticationMethod.KERBEROS;
  }
  SecurityUtil.setAuthenticationMethod(authMethod, conf);
  UserGroupInformation.setConfiguration(conf);
  InlineDispatcher rmDispatcher = new InlineDispatcher();

  ContainerAllocationExpirer containerAllocationExpirer =
      mock(ContainerAllocationExpirer.class);
  amLivelinessMonitor = mock(AMLivelinessMonitor.class);
  amFinishingMonitor = mock(AMLivelinessMonitor.class);
  writer = mock(RMApplicationHistoryWriter.class);
  MasterKeyData masterKeyData = amRMTokenManager.createNewMasterKey();
  when(amRMTokenManager.getMasterKey()).thenReturn(masterKeyData);
  rmContext =
      new RMContextImpl(rmDispatcher,
        containerAllocationExpirer, amLivelinessMonitor, amFinishingMonitor,
        null, amRMTokenManager,
        new RMContainerTokenSecretManager(conf),
        nmTokenManager,
        clientToAMTokenManager,
        writer);
  
  store = mock(RMStateStore.class);
  ((RMContextImpl) rmContext).setStateStore(store);
  publisher = mock(SystemMetricsPublisher.class);
  ((RMContextImpl) rmContext).setSystemMetricsPublisher(publisher);
  
  scheduler = mock(YarnScheduler.class);
  masterService = mock(ApplicationMasterService.class);
  applicationMasterLauncher = mock(ApplicationMasterLauncher.class);
  
  rmDispatcher.register(RMAppAttemptEventType.class,
      new TestApplicationAttemptEventDispatcher());

  rmDispatcher.register(RMAppEventType.class,
      new TestApplicationEventDispatcher());
  
  rmDispatcher.register(SchedulerEventType.class, 
      new TestSchedulerEventDispatcher());
  
  rmDispatcher.register(AMLauncherEventType.class, 
      new TestAMLauncherEventDispatcher());

  rmnodeEventHandler = mock(RMNodeImpl.class);
  rmDispatcher.register(RMNodeEventType.class, rmnodeEventHandler);

  rmDispatcher.init(conf);
  rmDispatcher.start();
  

  ApplicationId applicationId = MockApps.newAppID(appId++);
  ApplicationAttemptId applicationAttemptId =
      ApplicationAttemptId.newInstance(applicationId, 0);

  resourceScheduler = mock(ResourceScheduler.class);

  ApplicationResourceUsageReport appResUsgRpt =
      mock(ApplicationResourceUsageReport.class);
  when(appResUsgRpt.getMemorySeconds()).thenReturn(0L);
  when(appResUsgRpt.getVcoreSeconds()).thenReturn(0L);
  when(resourceScheduler
      .getAppResourceUsageReport((ApplicationAttemptId)Matchers.any()))
   .thenReturn(appResUsgRpt);
  spyRMContext = spy(rmContext);
  Mockito.doReturn(resourceScheduler).when(spyRMContext).getScheduler();


  final String user = MockApps.newUserName();
  final String queue = MockApps.newQueue();
  submissionContext = mock(ApplicationSubmissionContext.class);
  when(submissionContext.getQueue()).thenReturn(queue);
  Resource resource = BuilderUtils.newResource(1536, 1);
  ContainerLaunchContext amContainerSpec =
      BuilderUtils.newContainerLaunchContext(null, null,
          null, null, null, null);
  when(submissionContext.getAMContainerSpec()).thenReturn(amContainerSpec);
  when(submissionContext.getResource()).thenReturn(resource);

  unmanagedAM = false;
  
  application = mock(RMAppImpl.class);
  applicationAttempt =
      new RMAppAttemptImpl(applicationAttemptId, spyRMContext, scheduler,
          masterService, submissionContext, new Configuration(), false,
          BuilderUtils.newResourceRequest(
              RMAppAttemptImpl.AM_CONTAINER_PRIORITY, ResourceRequest.ANY,
              submissionContext.getResource(), 1));

  when(application.getCurrentAppAttempt()).thenReturn(applicationAttempt);
  when(application.getApplicationId()).thenReturn(applicationId);
  spyRMContext.getRMApps().put(application.getApplicationId(), application);

  testAppAttemptNewState();
}
 

protected SystemMetricsPublisher createSystemMetricsPublisher() {
  return new SystemMetricsPublisher(); 
}
 

@Test(timeout=2000)
public void testUpdateResourceOnNode() throws Exception {
  AsyncDispatcher dispatcher = new InlineDispatcher();
  Configuration conf = new Configuration();
  RMContainerTokenSecretManager containerTokenSecretManager =
      new RMContainerTokenSecretManager(conf);
  containerTokenSecretManager.rollMasterKey();
  NMTokenSecretManagerInRM nmTokenSecretManager =
      new NMTokenSecretManagerInRM(conf);
  nmTokenSecretManager.rollMasterKey();
  RMApplicationHistoryWriter writer = mock(RMApplicationHistoryWriter.class);
  
  FifoScheduler scheduler = new FifoScheduler(){
    @SuppressWarnings("unused")
    public Map<NodeId, FiCaSchedulerNode> getNodes(){
      return nodes;
    }
  };
  RMContext rmContext = new RMContextImpl(dispatcher, null, null, null, null,
      null, containerTokenSecretManager, nmTokenSecretManager, null, writer,
      scheduler);
  ((RMContextImpl) rmContext).setSystemMetricsPublisher(
      mock(SystemMetricsPublisher.class));

  scheduler.setRMContext(rmContext);
  scheduler.init(conf);
  scheduler.start();
  scheduler.reinitialize(new Configuration(), rmContext);
  RMNode node0 = MockNodes.newNodeInfo(1,
      Resources.createResource(2048, 4), 1, "127.0.0.1");
  NodeAddedSchedulerEvent nodeEvent1 = new NodeAddedSchedulerEvent(node0);
  scheduler.handle(nodeEvent1);
  
  Method method = scheduler.getClass().getDeclaredMethod("getNodes");
  @SuppressWarnings("unchecked")
  Map<NodeId, FiCaSchedulerNode> schedulerNodes = 
      (Map<NodeId, FiCaSchedulerNode>) method.invoke(scheduler);
  assertEquals(schedulerNodes.values().size(), 1);
  
  Resource newResource = Resources.createResource(1024, 4);
  
  NodeResourceUpdateSchedulerEvent node0ResourceUpdate = new 
      NodeResourceUpdateSchedulerEvent(node0, ResourceOption.newInstance(
          newResource, RMNode.OVER_COMMIT_TIMEOUT_MILLIS_DEFAULT));
  scheduler.handle(node0ResourceUpdate);
  
  // SchedulerNode's total resource and available resource are changed.
  assertEquals(schedulerNodes.get(node0.getNodeID()).getTotalResource()
      .getMemory(), 1024);
  assertEquals(schedulerNodes.get(node0.getNodeID()).
      getAvailableResource().getMemory(), 1024);
  QueueInfo queueInfo = scheduler.getQueueInfo(null, false, false);
  Assert.assertEquals(0.0f, queueInfo.getCurrentCapacity(), 0.0f);
  
  int _appId = 1;
  int _appAttemptId = 1;
  ApplicationAttemptId appAttemptId = createAppAttemptId(_appId,
      _appAttemptId);
  createMockRMApp(appAttemptId, rmContext);

  AppAddedSchedulerEvent appEvent =
      new AppAddedSchedulerEvent(appAttemptId.getApplicationId(), "queue1",
        "user1");
  scheduler.handle(appEvent);
  AppAttemptAddedSchedulerEvent attemptEvent =
      new AppAttemptAddedSchedulerEvent(appAttemptId, false);
  scheduler.handle(attemptEvent);

  int memory = 1024;
  int priority = 1;

  List<ResourceRequest> ask = new ArrayList<ResourceRequest>();
  ResourceRequest nodeLocal = createResourceRequest(memory,
      node0.getHostName(), priority, 1);
  ResourceRequest rackLocal = createResourceRequest(memory,
      node0.getRackName(), priority, 1);
  ResourceRequest any = createResourceRequest(memory, ResourceRequest.ANY, priority,
      1);
  ask.add(nodeLocal);
  ask.add(rackLocal);
  ask.add(any);
  scheduler.allocate(appAttemptId, ask, new ArrayList<ContainerId>(), null, null);

  // Before the node update event, there are one local request
  Assert.assertEquals(1, nodeLocal.getNumContainers());

  NodeUpdateSchedulerEvent node0Update = new NodeUpdateSchedulerEvent(node0);
  // Now schedule.
  scheduler.handle(node0Update);

  // After the node update event, check no local request
  Assert.assertEquals(0, nodeLocal.getNumContainers());
  // Also check that one container was scheduled
  SchedulerAppReport info = scheduler.getSchedulerAppInfo(appAttemptId);
  Assert.assertEquals(1, info.getLiveContainers().size());
  // And check the default Queue now is full.
  queueInfo = scheduler.getQueueInfo(null, false, false);
  Assert.assertEquals(1.0f, queueInfo.getCurrentCapacity(), 0.0f);
}
 

@Test
public void testGetAppToUnreserve() throws Exception {

  CapacitySchedulerConfiguration csConf = new CapacitySchedulerConfiguration();
  setup(csConf);
  final String user_0 = "user_0";
  final ApplicationAttemptId appAttemptId_0 = TestUtils
      .getMockApplicationAttemptId(0, 0);
  LeafQueue a = stubLeafQueue((LeafQueue) queues.get(A));
  FiCaSchedulerApp app_0 = new FiCaSchedulerApp(appAttemptId_0, user_0, a,
      mock(ActiveUsersManager.class), spyRMContext);

  String host_0 = "host_0";
  FiCaSchedulerNode node_0 = TestUtils.getMockNode(host_0, DEFAULT_RACK, 0,
      8 * GB);
  String host_1 = "host_1";
  FiCaSchedulerNode node_1 = TestUtils.getMockNode(host_1, DEFAULT_RACK, 0,
      8 * GB);
  
  Resource clusterResource = Resources.createResource(2 * 8 * GB);

  // Setup resource-requests
  Priority priorityMap = TestUtils.createMockPriority(5);
  Resource capability = Resources.createResource(2*GB, 0);

  RMApplicationHistoryWriter writer = mock(RMApplicationHistoryWriter.class);
  SystemMetricsPublisher publisher = mock(SystemMetricsPublisher.class);
  RMContext rmContext = mock(RMContext.class);
  ContainerAllocationExpirer expirer =
    mock(ContainerAllocationExpirer.class);
  DrainDispatcher drainDispatcher = new DrainDispatcher();
  when(rmContext.getContainerAllocationExpirer()).thenReturn(expirer);
  when(rmContext.getDispatcher()).thenReturn(drainDispatcher);
  when(rmContext.getRMApplicationHistoryWriter()).thenReturn(writer);
  when(rmContext.getSystemMetricsPublisher()).thenReturn(publisher);
  ApplicationAttemptId appAttemptId = BuilderUtils.newApplicationAttemptId(
      app_0.getApplicationId(), 1);
  ContainerId containerId = BuilderUtils.newContainerId(appAttemptId, 1);
  Container container = TestUtils.getMockContainer(containerId,
      node_1.getNodeID(), Resources.createResource(2*GB), priorityMap);
  RMContainer rmContainer = new RMContainerImpl(container, appAttemptId,
      node_1.getNodeID(), "user", rmContext);

  Container container_1 = TestUtils.getMockContainer(containerId,
      node_0.getNodeID(), Resources.createResource(1*GB), priorityMap);
  RMContainer rmContainer_1 = new RMContainerImpl(container_1, appAttemptId,
      node_0.getNodeID(), "user", rmContext);

  // no reserved containers
  NodeId unreserveId =
      app_0.getNodeIdToUnreserve(priorityMap, capability,
          cs.getResourceCalculator(), clusterResource);
  assertEquals(null, unreserveId);

  // no reserved containers - reserve then unreserve
  app_0.reserve(node_0, priorityMap, rmContainer_1, container_1);
  app_0.unreserve(node_0, priorityMap);
  unreserveId = app_0.getNodeIdToUnreserve(priorityMap, capability,
      cs.getResourceCalculator(), clusterResource);
  assertEquals(null, unreserveId);

  // no container large enough is reserved
  app_0.reserve(node_0, priorityMap, rmContainer_1, container_1);
  unreserveId = app_0.getNodeIdToUnreserve(priorityMap, capability,
      cs.getResourceCalculator(), clusterResource);
  assertEquals(null, unreserveId);

  // reserve one that is now large enough
  app_0.reserve(node_1, priorityMap, rmContainer, container);
  unreserveId = app_0.getNodeIdToUnreserve(priorityMap, capability,
      cs.getResourceCalculator(), clusterResource);
  assertEquals(node_1.getNodeID(), unreserveId);
}
 

@Test
public void testFindNodeToUnreserve() throws Exception {

  CapacitySchedulerConfiguration csConf = new CapacitySchedulerConfiguration();
  setup(csConf);
  final String user_0 = "user_0";
  final ApplicationAttemptId appAttemptId_0 = TestUtils
      .getMockApplicationAttemptId(0, 0);
  LeafQueue a = stubLeafQueue((LeafQueue) queues.get(A));
  FiCaSchedulerApp app_0 = new FiCaSchedulerApp(appAttemptId_0, user_0, a,
      mock(ActiveUsersManager.class), spyRMContext);

  String host_1 = "host_1";
  FiCaSchedulerNode node_1 = TestUtils.getMockNode(host_1, DEFAULT_RACK, 0,
      8 * GB);

  // Setup resource-requests
  Priority priorityMap = TestUtils.createMockPriority(5);
  Resource capability = Resources.createResource(2 * GB, 0);

  RMApplicationHistoryWriter writer = mock(RMApplicationHistoryWriter.class);
  SystemMetricsPublisher publisher = mock(SystemMetricsPublisher.class);
  RMContext rmContext = mock(RMContext.class);
  ContainerAllocationExpirer expirer =
    mock(ContainerAllocationExpirer.class);
  DrainDispatcher drainDispatcher = new DrainDispatcher();
  when(rmContext.getContainerAllocationExpirer()).thenReturn(expirer);
  when(rmContext.getDispatcher()).thenReturn(drainDispatcher);
  when(rmContext.getRMApplicationHistoryWriter()).thenReturn(writer);
  when(rmContext.getSystemMetricsPublisher()).thenReturn(publisher);
  ApplicationAttemptId appAttemptId = BuilderUtils.newApplicationAttemptId(
      app_0.getApplicationId(), 1);
  ContainerId containerId = BuilderUtils.newContainerId(appAttemptId, 1);
  Container container = TestUtils.getMockContainer(containerId,
      node_1.getNodeID(), Resources.createResource(2*GB), priorityMap);
  RMContainer rmContainer = new RMContainerImpl(container, appAttemptId,
      node_1.getNodeID(), "user", rmContext);

  // nothing reserved
  boolean res = a.findNodeToUnreserve(csContext.getClusterResource(),
      node_1, app_0, priorityMap, capability);
  assertFalse(res);

  // reserved but scheduler doesn't know about that node.
  app_0.reserve(node_1, priorityMap, rmContainer, container);
  node_1.reserveResource(app_0, priorityMap, rmContainer);
  res = a.findNodeToUnreserve(csContext.getClusterResource(), node_1, app_0,
      priorityMap, capability);
  assertFalse(res);
}
 

@Override
public void setSystemMetricsPublisher(SystemMetricsPublisher systemMetricsPublisher) {
  this.systemMetricsPublisher = systemMetricsPublisher;
}
 

@Override
public SystemMetricsPublisher getSystemMetricsPublisher() {
  return systemMetricsPublisher;
}
 

@Override
@Before
public void setUp() throws Exception {
  super.setUp();
  this.baseStateStoreDirectory = "/tmp/nm-heartbeat-handler-test";
  context = new MockRMContext();
  context.setDispatcher(TestObjectFactory.getMockDispatcher());
  context.setSystemMetricsPublisher(new SystemMetricsPublisher());
  
  profileZero = TestObjectFactory.getServiceResourceProfile("zero", Double.valueOf(0.0), Double.valueOf(0.0), 
           Long.valueOf(0), Long.valueOf(0));
  profileSmall = TestObjectFactory.getServiceResourceProfile("small", Double.valueOf(2.0), Double.valueOf(2048.0), 
          Long.valueOf(1), Long.valueOf(1024));
  
  nodeOne = TestObjectFactory.getRMNode("localhost-one", 8800, Resource.newInstance(0, 0));
  nodeTwo = TestObjectFactory.getRMNode("localhost-two", 8800, Resource.newInstance(1024, 2));

  sNodeOne = new FSSchedulerNode(nodeOne, false);
  sNodeTwo = new FSSchedulerNode(nodeTwo, false);
  nodeTaskOne = TestObjectFactory.getNodeTask("localhost-one", profileZero);
  nodeTaskTwo = TestObjectFactory.getNodeTask("localhost-two", profileSmall);
  
  ConcurrentMap<NodeId, RMNode> rmNodes = new ConcurrentHashMap<NodeId, RMNode>();
  rmNodes.put(nodeOne.getNodeID(), nodeOne);
  rmNodes.put(nodeTwo.getNodeID(), nodeTwo);
  context.setRMNodes(rmNodes);
  
  store = new NodeStore();
  store.add(sNodeOne);
  store.add(sNodeTwo);
  
  MyriadDriver driver = TestObjectFactory.getMyriadDriver(new MockSchedulerDriver());
  olManager = new OfferLifecycleManager(store, driver);
  
  state = TestObjectFactory.getSchedulerState(new MyriadConfiguration(), "/tmp/nm-heartbeat-handler-test");
  state.addNodes(Lists.newArrayList(nodeTaskOne, nodeTaskTwo));
  MyriadFairScheduler scheduler = TestObjectFactory.getMyriadFairScheduler(context);
 
  scheduler.addNode(sNodeOne);
  scheduler.addNode(sNodeTwo);
  
  manager = new YarnNodeCapacityManager(new CompositeInterceptor(), scheduler, 
          context, driver, olManager, store, state, new TaskUtils(this.cfg));
  handler = new NMHeartBeatHandler(new CompositeInterceptor(), scheduler, 
          driver, manager, olManager, store, state, cfg.getNodeManagerConfiguration());
}
 

@Before
public void setUp() throws Exception {
  conf = new YarnConfiguration();
  AuthenticationMethod authMethod = AuthenticationMethod.SIMPLE;
  if (isSecurityEnabled) {
    authMethod = AuthenticationMethod.KERBEROS;
  }
  SecurityUtil.setAuthenticationMethod(authMethod, conf);
  UserGroupInformation.setConfiguration(conf);

  rmDispatcher = new DrainDispatcher();
  ContainerAllocationExpirer containerAllocationExpirer = 
      mock(ContainerAllocationExpirer.class);
  AMLivelinessMonitor amLivelinessMonitor = mock(AMLivelinessMonitor.class);
  AMLivelinessMonitor amFinishingMonitor = mock(AMLivelinessMonitor.class);
  store = mock(RMStateStore.class);
  writer = mock(RMApplicationHistoryWriter.class);
  DelegationTokenRenewer renewer = mock(DelegationTokenRenewer.class);
  RMContext realRMContext = 
      new RMContextImpl(rmDispatcher,
        containerAllocationExpirer, amLivelinessMonitor, amFinishingMonitor,
        renewer, new AMRMTokenSecretManager(conf, this.rmContext),
        new RMContainerTokenSecretManager(conf),
        new NMTokenSecretManagerInRM(conf),
        new ClientToAMTokenSecretManagerInRM(),
        writer);
  ((RMContextImpl)realRMContext).setStateStore(store);
  publisher = mock(SystemMetricsPublisher.class);
  ((RMContextImpl)realRMContext).setSystemMetricsPublisher(publisher);

  this.rmContext = spy(realRMContext);

  ResourceScheduler resourceScheduler = mock(ResourceScheduler.class);
  doReturn(null).when(resourceScheduler)
            .getAppResourceUsageReport((ApplicationAttemptId)Matchers.any());
  doReturn(resourceScheduler).when(rmContext).getScheduler();

  rmDispatcher.register(RMAppAttemptEventType.class,
      new TestApplicationAttemptEventDispatcher(this.rmContext));

  rmDispatcher.register(RMAppEventType.class,
      new TestApplicationEventDispatcher(rmContext));
  
  rmDispatcher.register(RMAppManagerEventType.class,
      new TestApplicationManagerEventDispatcher());
  
  schedulerDispatcher = new TestSchedulerEventDispatcher();
  rmDispatcher.register(SchedulerEventType.class,
      schedulerDispatcher);
  
  rmDispatcher.init(conf);
  rmDispatcher.start();
}
 

@Private
@Unstable
public void setSystemMetricsPublisher(
    SystemMetricsPublisher systemMetricsPublisher) {
  this.systemMetricsPublisher = systemMetricsPublisher;
}
 

@Private
@Unstable
public SystemMetricsPublisher getSystemMetricsPublisher() {
  return systemMetricsPublisher;
}
 

protected SystemMetricsPublisher createSystemMetricsPublisher() {
  return new SystemMetricsPublisher(); 
}
 

@Override
public void setSystemMetricsPublisher(
    SystemMetricsPublisher systemMetricsPublisher) {
  activeServiceContext.setSystemMetricsPublisher(systemMetricsPublisher);
}
 

@Override
public SystemMetricsPublisher getSystemMetricsPublisher() {
  return activeServiceContext.getSystemMetricsPublisher();
}
 

@Test
public void testReleaseWhileRunning() {

  DrainDispatcher drainDispatcher = new DrainDispatcher();
  EventHandler<RMAppAttemptEvent> appAttemptEventHandler = mock(EventHandler.class);
  EventHandler generic = mock(EventHandler.class);
  drainDispatcher.register(RMAppAttemptEventType.class,
      appAttemptEventHandler);
  drainDispatcher.register(RMNodeEventType.class, generic);
  drainDispatcher.init(new YarnConfiguration());
  drainDispatcher.start();
  NodeId nodeId = BuilderUtils.newNodeId("host", 3425);
  ApplicationId appId = BuilderUtils.newApplicationId(1, 1);
  ApplicationAttemptId appAttemptId = BuilderUtils.newApplicationAttemptId(
      appId, 1);
  ContainerId containerId = BuilderUtils.newContainerId(appAttemptId, 1);
  ContainerAllocationExpirer expirer = mock(ContainerAllocationExpirer.class);

  Resource resource = BuilderUtils.newResource(512, 1, 1);
  Priority priority = BuilderUtils.newPriority(5);

  Container container = BuilderUtils.newContainer(containerId, nodeId,
      "host:3465", resource, priority, null);
  ConcurrentMap<ApplicationId, RMApp> rmApps =
      spy(new ConcurrentHashMap<ApplicationId, RMApp>());
  RMApp rmApp = mock(RMApp.class);
  when(rmApp.getRMAppAttempt((ApplicationAttemptId)Matchers.any())).thenReturn(null);
  Mockito.doReturn(rmApp).when(rmApps).get((ApplicationId)Matchers.any());

  RMApplicationHistoryWriter writer = mock(RMApplicationHistoryWriter.class);
  SystemMetricsPublisher publisher = mock(SystemMetricsPublisher.class);
  RMContext rmContext = mock(RMContext.class);
  when(rmContext.getDispatcher()).thenReturn(drainDispatcher);
  when(rmContext.getContainerAllocationExpirer()).thenReturn(expirer);
  when(rmContext.getRMApplicationHistoryWriter()).thenReturn(writer);
  when(rmContext.getRMApps()).thenReturn(rmApps);
  when(rmContext.getSystemMetricsPublisher()).thenReturn(publisher);
  when(rmContext.getYarnConfiguration()).thenReturn(new YarnConfiguration());
  RMContainer rmContainer = new RMContainerImpl(container, appAttemptId,
      nodeId, "user", rmContext);

  assertEquals(RMContainerState.NEW, rmContainer.getState());
  assertEquals(resource, rmContainer.getAllocatedResource());
  assertEquals(nodeId, rmContainer.getAllocatedNode());
  assertEquals(priority, rmContainer.getAllocatedPriority());
  verify(writer).containerStarted(any(RMContainer.class));
  verify(publisher).containerCreated(any(RMContainer.class), anyLong());

  rmContainer.handle(new RMContainerEvent(containerId,
      RMContainerEventType.START));
  drainDispatcher.await();
  assertEquals(RMContainerState.ALLOCATED, rmContainer.getState());
  rmContainer.handle(new RMContainerEvent(containerId,
      RMContainerEventType.ACQUIRED));
  drainDispatcher.await();
  assertEquals(RMContainerState.ACQUIRED, rmContainer.getState());

  rmContainer.handle(new RMContainerEvent(containerId,
      RMContainerEventType.LAUNCHED));
  drainDispatcher.await();
  assertEquals(RMContainerState.RUNNING, rmContainer.getState());
  assertEquals("http://host:3465/node/containerlogs/container_1_0001_01_000001/user",
      rmContainer.getLogURL());

  // In RUNNING state. Verify RELEASED and associated actions.
  reset(appAttemptEventHandler);
  ContainerStatus containerStatus = SchedulerUtils
      .createAbnormalContainerStatus(containerId,
          SchedulerUtils.RELEASED_CONTAINER);
  rmContainer.handle(new RMContainerFinishedEvent(containerId,
      containerStatus, RMContainerEventType.RELEASED));
  drainDispatcher.await();
  assertEquals(RMContainerState.RELEASED, rmContainer.getState());
  assertEquals(SchedulerUtils.RELEASED_CONTAINER,
      rmContainer.getDiagnosticsInfo());
  assertEquals(ContainerExitStatus.ABORTED,
      rmContainer.getContainerExitStatus());
  assertEquals(ContainerState.COMPLETE, rmContainer.getContainerState());
  verify(writer).containerFinished(any(RMContainer.class));
  verify(publisher).containerFinished(any(RMContainer.class), anyLong());

  ArgumentCaptor<RMAppAttemptContainerFinishedEvent> captor = ArgumentCaptor
      .forClass(RMAppAttemptContainerFinishedEvent.class);
  verify(appAttemptEventHandler).handle(captor.capture());
  RMAppAttemptContainerFinishedEvent cfEvent = captor.getValue();
  assertEquals(appAttemptId, cfEvent.getApplicationAttemptId());
  assertEquals(containerStatus, cfEvent.getContainerStatus());
  assertEquals(RMAppAttemptEventType.CONTAINER_FINISHED, cfEvent.getType());
  
  // In RELEASED state. A FINIHSED event may come in.
  rmContainer.handle(new RMContainerFinishedEvent(containerId, SchedulerUtils
      .createAbnormalContainerStatus(containerId, "FinishedContainer"),
      RMContainerEventType.FINISHED));
  assertEquals(RMContainerState.RELEASED, rmContainer.getState());
}
 

@Test
public void testExpireWhileRunning() {

  DrainDispatcher drainDispatcher = new DrainDispatcher();
  EventHandler<RMAppAttemptEvent> appAttemptEventHandler = mock(EventHandler.class);
  EventHandler generic = mock(EventHandler.class);
  drainDispatcher.register(RMAppAttemptEventType.class,
      appAttemptEventHandler);
  drainDispatcher.register(RMNodeEventType.class, generic);
  drainDispatcher.init(new YarnConfiguration());
  drainDispatcher.start();
  NodeId nodeId = BuilderUtils.newNodeId("host", 3425);
  ApplicationId appId = BuilderUtils.newApplicationId(1, 1);
  ApplicationAttemptId appAttemptId = BuilderUtils.newApplicationAttemptId(
      appId, 1);
  ContainerId containerId = BuilderUtils.newContainerId(appAttemptId, 1);
  ContainerAllocationExpirer expirer = mock(ContainerAllocationExpirer.class);

  Resource resource = BuilderUtils.newResource(512, 1, 1);
  Priority priority = BuilderUtils.newPriority(5);

  Container container = BuilderUtils.newContainer(containerId, nodeId,
      "host:3465", resource, priority, null);

  RMApplicationHistoryWriter writer = mock(RMApplicationHistoryWriter.class);
  SystemMetricsPublisher publisher = mock(SystemMetricsPublisher.class);
  RMContext rmContext = mock(RMContext.class);
  when(rmContext.getDispatcher()).thenReturn(drainDispatcher);
  when(rmContext.getContainerAllocationExpirer()).thenReturn(expirer);
  when(rmContext.getRMApplicationHistoryWriter()).thenReturn(writer);
  when(rmContext.getSystemMetricsPublisher()).thenReturn(publisher);
  when(rmContext.getYarnConfiguration()).thenReturn(new YarnConfiguration());
  RMContainer rmContainer = new RMContainerImpl(container, appAttemptId,
      nodeId, "user", rmContext);

  assertEquals(RMContainerState.NEW, rmContainer.getState());
  assertEquals(resource, rmContainer.getAllocatedResource());
  assertEquals(nodeId, rmContainer.getAllocatedNode());
  assertEquals(priority, rmContainer.getAllocatedPriority());
  verify(writer).containerStarted(any(RMContainer.class));
  verify(publisher).containerCreated(any(RMContainer.class), anyLong());

  rmContainer.handle(new RMContainerEvent(containerId,
      RMContainerEventType.START));
  drainDispatcher.await();
  assertEquals(RMContainerState.ALLOCATED, rmContainer.getState());

  rmContainer.handle(new RMContainerEvent(containerId,
      RMContainerEventType.ACQUIRED));
  drainDispatcher.await();
  assertEquals(RMContainerState.ACQUIRED, rmContainer.getState());

  rmContainer.handle(new RMContainerEvent(containerId,
      RMContainerEventType.LAUNCHED));
  drainDispatcher.await();
  assertEquals(RMContainerState.RUNNING, rmContainer.getState());
  assertEquals("http://host:3465/node/containerlogs/container_1_0001_01_000001/user",
      rmContainer.getLogURL());

  // In RUNNING state. Verify EXPIRE and associated actions.
  reset(appAttemptEventHandler);
  ContainerStatus containerStatus = SchedulerUtils
      .createAbnormalContainerStatus(containerId,
          SchedulerUtils.EXPIRED_CONTAINER);
  rmContainer.handle(new RMContainerFinishedEvent(containerId,
      containerStatus, RMContainerEventType.EXPIRE));
  drainDispatcher.await();
  assertEquals(RMContainerState.RUNNING, rmContainer.getState());
  verify(writer, never()).containerFinished(any(RMContainer.class));
  verify(publisher, never()).containerFinished(any(RMContainer.class),
      anyLong());
}