@Test
  public void testRequestFailureOnStageFailure() throws Exception {
    Clusters fsm = mock(Clusters.class);
    Cluster oneClusterMock = mock(Cluster.class);
    Service serviceObj = mock(Service.class);
    ServiceComponent scomp = mock(ServiceComponent.class);
    ServiceComponentHost sch = mock(ServiceComponentHost.class);
    UnitOfWork unitOfWork = mock(UnitOfWork.class);
    AmbariEventPublisher ambariEventPublisher = mock(AmbariEventPublisher.class);
    RequestFactory requestFactory = mock(RequestFactory.class);
    when(fsm.getCluster(anyString())).thenReturn(oneClusterMock);
    when(oneClusterMock.getService(anyString())).thenReturn(serviceObj);
    when(serviceObj.getServiceComponent(anyString())).thenReturn(scomp);
    when(scomp.getServiceComponentHost(anyString())).thenReturn(sch);
    when(serviceObj.getCluster()).thenReturn(oneClusterMock);

    HashMap<String, ServiceComponentHost> hosts =
        new HashMap<>();
    hosts.put(hostname, sch);
    when(scomp.getServiceComponentHosts()).thenReturn(hosts);

    final List<Stage> stages = new ArrayList<>();

    stages.add(
        getStageWithSingleTask(
            hostname, "cluster1", Role.NAMENODE, RoleCommand.UPGRADE, Service.Type.HDFS, 1, 1, 1));

    List<Stage> firstStageInProgress = Collections.singletonList(stages.get(0));

    stages.add(
        getStageWithSingleTask(
            hostname, "cluster1", Role.DATANODE, RoleCommand.UPGRADE, Service.Type.HDFS, 2, 2, 1));

    Host host = mock(Host.class);
    when(fsm.getHost(anyString())).thenReturn(host);
    when(host.getState()).thenReturn(HostState.HEALTHY);
    when(host.getHostName()).thenReturn(hostname);

    ActionDBAccessor db = mock(ActionDBAccessor.class);

    RequestEntity request = mock(RequestEntity.class);
    when(request.isExclusive()).thenReturn(false);
    when(request.getClusterHostInfo()).thenReturn(CLUSTER_HOST_INFO);
    when(db.getRequestEntity(anyLong())).thenReturn(request);

    when(db.getCommandsInProgressCount()).thenReturn(stages.size());
    when(db.getFirstStageInProgressPerRequest()).thenReturn(firstStageInProgress);
    doAnswer(new Answer<Void>() {
      @Override
      public Void answer(InvocationOnMock invocation) throws Throwable {
        List<CommandReport> reports = (List<CommandReport>) invocation.getArguments()[0];
        for (CommandReport report : reports) {
          String actionId = report.getActionId();
          long[] requestStageIds = StageUtils.getRequestStage(actionId);
          Long requestId = requestStageIds[0];
          Long stageId = requestStageIds[1];
          Long id = report.getTaskId();
          for (Stage stage : stages) {
            if (requestId.equals(stage.getRequestId()) && stageId.equals(stage.getStageId())) {
              for (HostRoleCommand hostRoleCommand : stage.getOrderedHostRoleCommands()) {
                if (hostRoleCommand.getTaskId() == id) {
                  hostRoleCommand.setStatus(HostRoleStatus.valueOf(report.getStatus()));
                }
              }
            }
          }

        }

        return null;
      }
    }).when(db).updateHostRoleStates(anyCollectionOf(CommandReport.class));

    when(db.getTask(anyLong())).thenAnswer(new Answer<Object>() {
      @Override
      public Object answer(InvocationOnMock invocation) throws Throwable {
        Long taskId = (Long) invocation.getArguments()[0];
        for (Stage stage : stages) {
          for (HostRoleCommand command : stage.getOrderedHostRoleCommands()) {
            if (taskId.equals(command.getTaskId())) {
              return command;
            }
          }
        }
        return null;
      }
    });
    doAnswer(new Answer<Collection<HostRoleCommandEntity>>() {
      @Override
      public Collection<HostRoleCommandEntity> answer(InvocationOnMock invocation) throws Throwable {
        Long requestId = (Long) invocation.getArguments()[0];
        List<HostRoleCommandEntity> abortedCommands = Lists.newArrayList();

        for (Stage stage : stages) {
          if (requestId.equals(stage.getRequestId())) {
            for (HostRoleCommand command : stage.getOrderedHostRoleCommands()) {
              if (command.getStatus() == HostRoleStatus.QUEUED ||
                  command.getStatus() == HostRoleStatus.IN_PROGRESS ||
                  command.getStatus() == HostRoleStatus.PENDING) {
                command.setStatus(HostRoleStatus.ABORTED);

                HostRoleCommandEntity hostRoleCommandEntity = command.constructNewPersistenceEntity();
                hostRoleCommandEntity.setStage(stage.constructNewPersistenceEntity());

                abortedCommands.add(hostRoleCommandEntity);
              }
            }
          }
        }

        return abortedCommands;
      }
    }).when(db).abortOperation(anyLong());

    Properties properties = new Properties();
    Configuration conf = new Configuration(properties);

    Capture<Collection<HostRoleCommand>> cancelCommandList = EasyMock.newCapture();
    ActionScheduler scheduler = EasyMock.createMockBuilder(ActionScheduler.class).
        withConstructor((long)100, (long)50, db, fsm, 3,
            new HostsMap((String) null),
            unitOfWork, EasyMock.createNiceMock(AmbariEventPublisher.class), conf,
            entityManagerProviderMock, mock(HostRoleCommandDAO.class),
            mock(HostRoleCommandFactory.class),
            mock(RoleCommandOrderProvider.class),
            mock(AgentCommandsPublisher.class)).
        addMockedMethod("cancelHostRoleCommands").
        createMock();
    scheduler.cancelHostRoleCommands(EasyMock.capture(cancelCommandList),
        EasyMock.eq(ActionScheduler.FAILED_TASK_ABORT_REASONING));
    EasyMock.expectLastCall().once();
    EasyMock.replay(scheduler);

    ActionManager am = new ActionManager(db, requestFactory, scheduler);

    scheduler.doWork();

    List<CommandReport> reports = new ArrayList<>();
    reports.add(getCommandReport(HostRoleStatus.FAILED, Role.NAMENODE, Service.Type.HDFS, "1-1", 1));
    am.processTaskResponse(hostname, reports, CommandUtils.convertToTaskIdCommandMap(stages.get(0).getOrderedHostRoleCommands()));

    scheduler.doWork();
    Assert.assertEquals(HostRoleStatus.FAILED, stages.get(0).getHostRoleStatus(hostname, "NAMENODE"));
    Assert.assertEquals(HostRoleStatus.ABORTED, stages.get(1).getHostRoleStatus(hostname, "DATANODE"));
    Assert.assertEquals(cancelCommandList.getValue().size(), 1);
    EasyMock.verify(scheduler, entityManagerProviderMock);
  }
Verifies that ActionScheduler allows to execute background tasks in parallel
ambari-server/src/test/java/org/apache/ambari/server/actionmanager/TestActionScheduler.java