@Test
@SuppressWarnings("unchecked")
public void testUpgradeSpecificHandling() throws Exception {

  Cluster cluster = heartbeatTestHelper.getDummyCluster();
  Service hdfs = cluster.addService(HDFS);
  hdfs.addServiceComponent(DATANODE);
  hdfs.getServiceComponent(DATANODE).addServiceComponentHost(DummyHostname1);
  ActionQueue aq = new ActionQueue();
  ServiceComponentHost serviceComponentHost1 = clusters.getCluster(DummyCluster).getService(HDFS).
      getServiceComponent(DATANODE).getServiceComponentHost(DummyHostname1);
  serviceComponentHost1.setState(State.UPGRADING);
  HeartBeat hb = new HeartBeat();
  hb.setTimestamp(System.currentTimeMillis());
  hb.setResponseId(0);
  hb.setHostname(DummyHostname1);
  hb.setNodeStatus(new HostStatus(HostStatus.Status.HEALTHY, DummyHostStatus));
  List<CommandReport> reports = new ArrayList<CommandReport>();
  CommandReport cr = new CommandReport();
  cr.setActionId(StageUtils.getActionId(requestId, stageId));
  cr.setTaskId(1);
  cr.setClusterName(DummyCluster);
  cr.setServiceName(HDFS);
  cr.setRole(DATANODE);
  cr.setRoleCommand("INSTALL");
  cr.setStatus(HostRoleStatus.IN_PROGRESS.toString());
  cr.setStdErr("none");
  cr.setStdOut("dummy output");
  cr.setExitCode(777);
  reports.add(cr);
  hb.setReports(reports);
  hb.setComponentStatus(new ArrayList<ComponentStatus>());
  final HostRoleCommand command = hostRoleCommandFactory.create(DummyHostname1,
      Role.DATANODE, null, null);
  ActionManager am = actionManagerTestHelper.getMockActionManager();
  expect(am.getTasks(EasyMock.<List<Long>>anyObject())).andReturn(
      new ArrayList<HostRoleCommand>() {{
        add(command);
      }}).anyTimes();
  replay(am);
  HeartBeatHandler handler = heartbeatTestHelper.getHeartBeatHandler(am, aq);
  HeartbeatProcessor heartbeatProcessor = handler.getHeartbeatProcessor();
  heartbeatProcessor.processHeartbeat(hb);
  assertEquals("Host state should  be " + State.UPGRADING,
      State.UPGRADING, serviceComponentHost1.getState());
  hb.setTimestamp(System.currentTimeMillis());
  hb.setResponseId(1);
  cr.setStatus(HostRoleStatus.COMPLETED.toString());
  cr.setExitCode(0);
  heartbeatProcessor.processHeartbeat(hb);
  assertEquals("Host state should be " + State.INSTALLED,
      State.INSTALLED, serviceComponentHost1.getState());
  serviceComponentHost1.setState(State.UPGRADING);
  hb.setTimestamp(System.currentTimeMillis());
  hb.setResponseId(2);
  cr.setStatus(HostRoleStatus.FAILED.toString());
  cr.setExitCode(3);
  heartbeatProcessor.processHeartbeat(hb);
  assertEquals("Host state should be " + State.UPGRADING,
      State.UPGRADING, serviceComponentHost1.getState());
  serviceComponentHost1.setState(State.UPGRADING);
  hb.setTimestamp(System.currentTimeMillis());
  hb.setResponseId(3);
  cr.setStatus(HostRoleStatus.PENDING.toString());
  cr.setExitCode(55);
  heartbeatProcessor.processHeartbeat(hb);
  assertEquals("Host state should be " + State.UPGRADING,
      State.UPGRADING, serviceComponentHost1.getState());
  serviceComponentHost1.setState(State.UPGRADING);
  hb.setTimestamp(System.currentTimeMillis());
  hb.setResponseId(4);
  cr.setStatus(HostRoleStatus.QUEUED.toString());
  cr.setExitCode(55);
  heartbeatProcessor.processHeartbeat(hb);
  assertEquals("Host state should be " + State.UPGRADING,
      State.UPGRADING, serviceComponentHost1.getState());
}