@Test
  public void testCancelRequests() 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);
    AgentCommandsPublisher agentCommandsPublisher = mock(AgentCommandsPublisher.class);
    HostRoleCommandDAO hostRoleCommandDAO = mock(HostRoleCommandDAO.class);
    HostRoleCommandFactory hostRoleCommandFactory = mock(HostRoleCommandFactory.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);

    final Long hostId = 1L;

    HostEntity hostEntity = new HostEntity();
    hostEntity.setHostName(hostname);
    hostEntity.setHostId(hostId);
    hostDAO.create(hostEntity);

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

    // Create a single request with 3 stages, each with a single task - the first stage will be completed and should not
    // be included when cancelling the unfinished tasks of the request
    long requestId = 1;
    final List<Stage> allStages = new ArrayList<>();
    final List<Stage> stagesInProgress = new ArrayList<>();
    final List<Stage> firstStageInProgress = new ArrayList<>();
    final List<HostRoleCommand> tasksInProgress = new ArrayList<>();
    final List<HostRoleCommandEntity> hrcEntitiesInProgress = new ArrayList<>();

    int secondaryNamenodeCmdTaskId = 1;
    int namenodeCmdTaskId = 2;
    int datanodeCmdTaskId = 3;

    Stage stageWithTask = getStageWithSingleTask(
        hostname, "cluster1", Role.SECONDARY_NAMENODE, RoleCommand.START,
        Service.Type.HDFS, secondaryNamenodeCmdTaskId, 1, (int) requestId);

    // complete the first stage
    stageWithTask.getOrderedHostRoleCommands().get(0).setStatus(HostRoleStatus.COMPLETED);
    allStages.add(stageWithTask);

    stageWithTask = getStageWithSingleTask(
        hostname, "cluster1", Role.NAMENODE, RoleCommand.START,
        Service.Type.HDFS, namenodeCmdTaskId, 2, (int) requestId);

    tasksInProgress.addAll(stageWithTask.getOrderedHostRoleCommands());
    firstStageInProgress.add(stageWithTask);
    stagesInProgress.add(stageWithTask);
    allStages.add(stageWithTask);

    stageWithTask = getStageWithSingleTask(
        hostname, "cluster1", Role.DATANODE, RoleCommand.START,
        Service.Type.HDFS, datanodeCmdTaskId, 3, (int) requestId);

    tasksInProgress.addAll(stageWithTask.getOrderedHostRoleCommands());
    stagesInProgress.add(stageWithTask);
    allStages.add(stageWithTask);

    // convert HRC to HRCEntity for the mock DAO to use
    for (HostRoleCommand hostRoleCommand : tasksInProgress) {
      HostRoleCommandEntity entity = mock(HostRoleCommandEntity.class);
      when(entity.getTaskId()).thenReturn(hostRoleCommand.getTaskId());
      when(entity.getStageId()).thenReturn(hostRoleCommand.getStageId());
      when(entity.getRequestId()).thenReturn(hostRoleCommand.getRequestId());
      when(entity.getHostId()).thenReturn(hostRoleCommand.getHostId());
      when(entity.getHostName()).thenReturn(hostRoleCommand.getHostName());
      when(entity.getRole()).thenReturn(hostRoleCommand.getRole());
      when(entity.getStatus()).thenReturn(hostRoleCommand.getStatus());
      when(entity.getRoleCommand()).thenReturn(hostRoleCommand.getRoleCommand());

      hrcEntitiesInProgress.add(entity);

      when(hostRoleCommandFactory.createExisting(entity)).thenReturn(hostRoleCommand);
    }

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

    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(stagesInProgress.size());
    when(db.getFirstStageInProgressPerRequest()).thenReturn(stagesInProgress);
    when(db.getStagesInProgressForRequest(requestId)).thenReturn(stagesInProgress);
    when(db.getAllStages(anyLong())).thenReturn(allStages);

    List<HostRoleCommand> requestTasks = new ArrayList<>();

    for (Stage stage : allStages) {
      requestTasks.addAll(stage.getOrderedHostRoleCommands());
    }

    when(db.getRequestTasks(anyLong())).thenReturn(requestTasks);

    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 : stagesInProgress) {
            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 : allStages) {
          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 : stagesInProgress) {
          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());
    Map<Long, List<AgentCommand>> commands = new HashMap<>();

    doAnswer(new Answer<Void>() {
      @Override
      public Void answer(InvocationOnMock invocation) throws Throwable {
        Long hostId = (Long) invocation.getArguments()[0];
        if (!commands.containsKey(hostId)) {
          commands.put(hostId, new ArrayList<>());
        }
        commands.get(hostId).add((AgentCommand) invocation.getArguments()[1]);
        return null;
      }
    }).when(agentCommandsPublisher).sendAgentCommand(anyLong(), any(AgentCommand.class));

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

    when(hostRoleCommandDAO.findByRequestIdAndStatuses(requestId,
        HostRoleStatus.NOT_COMPLETED_STATUSES)).thenReturn(hrcEntitiesInProgress);

    ActionScheduler scheduler = new ActionScheduler(100, 50, db, fsm, 3,
        new HostsMap((String) null), unitOfWork, null, conf, entityManagerProviderMock,
        hostRoleCommandDAO, hostRoleCommandFactory, agentCommandsPublisher);

    scheduler.doWork();

    String reason = "Some reason";

    scheduler.scheduleCancellingRequest(requestId, reason);

    scheduler.doWork();
    Assert.assertEquals(HostRoleStatus.COMPLETED, allStages.get(0).getHostRoleStatus(hostname, "SECONDARY_NAMENODE"));
    Assert.assertEquals(HostRoleStatus.ABORTED, allStages.get(1).getHostRoleStatus(hostname, "NAMENODE"));
    Assert.assertEquals(HostRoleStatus.ABORTED, allStages.get(2).getHostRoleStatus(hostname, "DATANODE"));

    Assert.assertEquals(1, commands.get(hostId).size()); // Cancel commands should be generated only for 1 stage

    CancelCommand cancelCommand = (CancelCommand) commands.get(hostId).get(0);
    Assert.assertEquals(cancelCommand.getTargetTaskId(), namenodeCmdTaskId);
    Assert.assertEquals(cancelCommand.getReason(), reason);
  }
Checks what happens when stage has an execution command for
host component that has been recently deleted
@throws Exception
ambari-server/src/test/java/org/apache/ambari/server/actionmanager/TestActionScheduler.java