/*
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you under the Apache License, Version 2.0 (the
 * "License"); you may not use this file except in compliance
 * with the License.  You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package org.apache.ambari.server.actionmanager;

import static org.easymock.EasyMock.createMock;
import static org.easymock.EasyMock.createNiceMock;
import static org.easymock.EasyMock.expect;
import static org.easymock.EasyMock.replay;
import static org.junit.Assert.assertEquals;
import static org.junit.Assert.assertFalse;
import static org.junit.Assert.assertTrue;
import static org.mockito.ArgumentMatchers.nullable;
import static org.mockito.Matchers.any;
import static org.mockito.Matchers.anyBoolean;
import static org.mockito.Matchers.anyCollectionOf;
import static org.mockito.Matchers.anyListOf;
import static org.mockito.Matchers.anyLong;
import static org.mockito.Matchers.anyString;
import static org.mockito.Matchers.eq;
import static org.mockito.Mockito.atLeastOnce;
import static org.mockito.Mockito.doAnswer;
import static org.mockito.Mockito.doReturn;
import static org.mockito.Mockito.mock;
import static org.mockito.Mockito.reset;
import static org.mockito.Mockito.spy;
import static org.mockito.Mockito.times;
import static org.mockito.Mockito.verify;
import static org.mockito.Mockito.when;

import java.lang.reflect.Method;
import java.lang.reflect.Type;
import java.sql.SQLException;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collection;
import java.util.Collections;
import java.util.HashMap;
import java.util.HashSet;
import java.util.List;
import java.util.Map;
import java.util.Properties;
import java.util.Set;
import java.util.TreeMap;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicInteger;

import jakarta.persistence.EntityManager;

import org.apache.ambari.server.AmbariException;
import org.apache.ambari.server.H2DatabaseCleaner;
import org.apache.ambari.server.Role;
import org.apache.ambari.server.RoleCommand;
import org.apache.ambari.server.ServiceComponentHostNotFoundException;
import org.apache.ambari.server.actionmanager.ActionScheduler.RoleStats;
import org.apache.ambari.server.agent.AgentCommand;
import org.apache.ambari.server.agent.AgentCommand.AgentCommandType;
import org.apache.ambari.server.agent.CancelCommand;
import org.apache.ambari.server.agent.CommandReport;
import org.apache.ambari.server.agent.ExecutionCommand;
import org.apache.ambari.server.configuration.Configuration;
import org.apache.ambari.server.controller.HostsMap;
import org.apache.ambari.server.events.AmbariEvent;
import org.apache.ambari.server.events.publishers.AgentCommandsPublisher;
import org.apache.ambari.server.events.publishers.AmbariEventPublisher;
import org.apache.ambari.server.metadata.RoleCommandOrder;
import org.apache.ambari.server.metadata.RoleCommandOrderProvider;
import org.apache.ambari.server.metadata.RoleCommandPair;
import org.apache.ambari.server.orm.GuiceJpaInitializer;
import org.apache.ambari.server.orm.InMemoryDefaultTestModule;
import org.apache.ambari.server.orm.dao.HostDAO;
import org.apache.ambari.server.orm.dao.HostRoleCommandDAO;
import org.apache.ambari.server.orm.entities.HostEntity;
import org.apache.ambari.server.orm.entities.HostRoleCommandEntity;
import org.apache.ambari.server.orm.entities.RequestEntity;
import org.apache.ambari.server.serveraction.MockServerAction;
import org.apache.ambari.server.serveraction.ServerActionExecutor;
import org.apache.ambari.server.state.Cluster;
import org.apache.ambari.server.state.Clusters;
import org.apache.ambari.server.state.Host;
import org.apache.ambari.server.state.HostState;
import org.apache.ambari.server.state.Service;
import org.apache.ambari.server.state.ServiceComponent;
import org.apache.ambari.server.state.ServiceComponentHost;
import org.apache.ambari.server.state.ServiceComponentHostEvent;
import org.apache.ambari.server.state.fsm.InvalidStateTransitionException;
import org.apache.ambari.server.state.svccomphost.ServiceComponentHostInstallEvent;
import org.apache.ambari.server.state.svccomphost.ServiceComponentHostOpFailedEvent;
import org.apache.ambari.server.state.svccomphost.ServiceComponentHostServerActionEvent;
import org.apache.ambari.server.state.svccomphost.ServiceComponentHostUpgradeEvent;
import org.apache.ambari.server.utils.CommandUtils;
import org.apache.ambari.server.utils.StageUtils;
import org.easymock.Capture;
import org.easymock.EasyMock;
import org.junit.After;
import org.junit.Assert;
import org.junit.Before;
import org.junit.Ignore;
import org.junit.Test;
import org.mockito.ArgumentCaptor;
import org.mockito.Mockito;
import org.mockito.invocation.InvocationOnMock;
import org.mockito.stubbing.Answer;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import com.google.common.collect.ArrayListMultimap;
import com.google.common.collect.Lists;
import com.google.common.collect.Multimap;
import com.google.common.reflect.TypeToken;
import com.google.common.util.concurrent.AtomicLongMap;
import com.google.inject.AbstractModule;
import com.google.inject.Guice;
import com.google.inject.Inject;
import com.google.inject.Injector;
import com.google.inject.Provider;
import com.google.inject.persist.UnitOfWork;

public class TestActionScheduler {

  private static final Logger log = LoggerFactory.getLogger(TestActionScheduler.class);
  private static final String CLUSTER_HOST_INFO = "{all_hosts=[c6403.ambari.apache.org," +
  		" c6401.ambari.apache.org, c6402.ambari.apache.org], slave_hosts=[c6403.ambari.apache.org," +
  		" c6401.ambari.apache.org, c6402.ambari.apache.org]}";
  private static final String CLUSTER_HOST_INFO_UPDATED = "{all_hosts=[c6401.ambari.apache.org,"
      + " c6402.ambari.apache.org], slave_hosts=[c6401.ambari.apache.org,"
      + " c6402.ambari.apache.org]}";

  private final Injector injector;

  private final String hostname = "ahost.ambari.apache.org";
  private final Long hostId = 1L;
  private final int MAX_CYCLE_ITERATIONS = 100;

  @Inject
  private HostRoleCommandFactory hostRoleCommandFactory;

  @Inject
  private StageFactory stageFactory;

  @Inject
  private HostDAO hostDAO;

  private Provider<EntityManager> entityManagerProviderMock = EasyMock.niceMock(Provider.class);


  public  TestActionScheduler(){
    injector = Guice.createInjector(new InMemoryDefaultTestModule());
  }

  @Before
  public void setup() throws Exception {
    injector.getInstance(GuiceJpaInitializer.class);
    injector.getInstance(StageUtils.class);
    injector.injectMembers(this);

    expect(entityManagerProviderMock.get()).andReturn(null);
    replay(entityManagerProviderMock);
  }

  @After
  public void teardown() throws AmbariException, SQLException {
    H2DatabaseCleaner.clearDatabaseAndStopPersistenceService(injector);
  }

  /**
   * This test sends a new action to the action scheduler and verifies that the action
   * shows up in the action queue.
   */
  @Test
  public void testActionSchedule() throws Exception {
    Type type = new TypeToken<Map<String, Set<String>>>() {}.getType();
    Map<String, Set<String>> clusterHostInfo = StageUtils.getGson().fromJson(CLUSTER_HOST_INFO, type);

    Properties properties = new Properties();
    Configuration conf = new Configuration(properties);
    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);
    when(fsm.getCluster(anyString())).thenReturn(oneClusterMock);
    when(fsm.getClusterById(anyLong())).thenReturn(oneClusterMock);
    when(oneClusterMock.getService(anyString())).thenReturn(serviceObj);
    when(oneClusterMock.getClusterId()).thenReturn(Long.valueOf(1L));
    when(serviceObj.getServiceComponent(anyString())).thenReturn(scomp);
    when(scomp.getServiceComponentHost(anyString())).thenReturn(sch);
    when(serviceObj.getCluster()).thenReturn(oneClusterMock);

    Host host = mock(Host.class);
    HashMap<String, ServiceComponentHost> hosts =
      new HashMap<>();
    hosts.put(hostname, sch);
    HostEntity hostEntity = new HostEntity();
    hostEntity.setHostName(hostname);
    hostEntity.setHostId(hostId);
    hostDAO.merge(hostEntity);
    when(scomp.getServiceComponentHosts()).thenReturn(hosts);

    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(ActionDBAccessorImpl.class);
    HostRoleCommandDAO hostRoleCommandDAOMock = mock(HostRoleCommandDAO.class);
    Mockito.doNothing().when(hostRoleCommandDAOMock).publishTaskCreateEvent(anyListOf(HostRoleCommand.class));
    Stage s = StageUtils.getATestStage(1, 977, hostname,
      "{\"host_param\":\"param_value\"}", "{\"stage_param\":\"param_value\"}");

    List<Stage> stages = Collections.singletonList(s);
    when(db.getCommandsInProgressCount()).thenReturn(stages.size());
    when(db.getFirstStageInProgressPerRequest()).thenReturn(stages);

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

    //Keep large number of attempts so that the task is not expired finally
    //Small action timeout to test rescheduling
    ActionScheduler scheduler = new ActionScheduler(100, 5, db, fsm,
        10000, new HostsMap((String) null), unitOfWork, null, conf, entityManagerProviderMock, hostRoleCommandDAOMock,
        null, agentCommandsPublisher);
    scheduler.setTaskTimeoutAdjustment(false);

    List<AgentCommand> commands = waitForQueueSize(hostId, agentCommandsPublisher, 1, scheduler);

    assertTrue(commands != null && commands.size() == 1);
    AgentCommand scheduledCommand = commands.get(0);
    assertTrue(scheduledCommand instanceof ExecutionCommand);
    assertEquals("1-977", ((ExecutionCommand) scheduledCommand).getCommandId());
    assertEquals(clusterHostInfo, ((ExecutionCommand) scheduledCommand).getClusterHostInfo());

    //The action status has not changed, it should be queued again.
    commands = waitForQueueSize(hostId, agentCommandsPublisher, 2, scheduler);

    // first command is cancel for previous
    assertTrue(commands != null && commands.size() == 2);
    scheduledCommand = commands.get(1);
    assertTrue(scheduledCommand instanceof ExecutionCommand);
    assertEquals("1-977", ((ExecutionCommand) scheduledCommand).getCommandId());
    assertEquals(clusterHostInfo, ((ExecutionCommand) scheduledCommand).getClusterHostInfo());

    //Now change the action status
    s.setHostRoleStatus(hostname, "NAMENODE", HostRoleStatus.COMPLETED);

    //Wait for sometime, it shouldn't be scheduled this time.
    scheduler.doWork();

    EasyMock.verify(entityManagerProviderMock);
  }

  private List<AgentCommand> waitForQueueSize(Long hostId, AgentCommandsPublisher agentCommandsPublisher,
                                              int expectedQueueSize, ActionScheduler scheduler) throws AmbariException {
    AtomicLongMap<Long> callCounterByHost = AtomicLongMap.create(new HashMap<Long, Long>());
    AtomicInteger methodCallCounter = new AtomicInteger(0);
    ArgumentCaptor<Multimap<Long, AgentCommand>> executionCommandCaptor = ArgumentCaptor.forClass((Class) Multimap.class);
    int counter = 0;
    resetAgentCommandPublisherMock(agentCommandsPublisher, callCounterByHost, methodCallCounter);
    while (counter++ <= MAX_CYCLE_ITERATIONS) {
      long callsCount = callCounterByHost.get(hostId);
      if (callsCount == expectedQueueSize) {
        verify(agentCommandsPublisher, times(methodCallCounter.get()))
            .sendAgentCommand(executionCommandCaptor.capture());
        Multimap<Long, AgentCommand> scheduledCommands = executionCommandCaptor.getValue();
        return new ArrayList<>(scheduledCommands.asMap().get(hostId));
      } else if (callsCount > expectedQueueSize) {
        Assert.fail("Expected size : " + expectedQueueSize + " Actual size="
            + counter);
      }
      try {
        scheduler.doWork();
      } catch (AmbariException e) {
        Assert.fail("Ambari exception : " + e.getMessage() + e.getStackTrace());
      }
    }
    return null;
  }

  private void resetAgentCommandPublisherMock(AgentCommandsPublisher agentCommandsPublisher,
                                              AtomicLongMap<Long> callCounterByHost,
                                              AtomicInteger methodCallCounter) throws AmbariException {
    reset(agentCommandsPublisher);
    Mockito.doAnswer((i) -> {
      Multimap<Long, AgentCommand> host = (Multimap<Long, AgentCommand>) i.getArguments()[0];
      if (host != null) {
        host.asMap().forEach((h, l) -> callCounterByHost.addAndGet(h, 1));
      }
      methodCallCounter.incrementAndGet();
      return null;
    }).when(agentCommandsPublisher).sendAgentCommand(any(Multimap.class));
  }


  /**
   * This test sends a new action to the action scheduler and verifies that the action
   * shows up in the action queue in case of DEPENDENCY_ORDERED execution type, with RoleCommand
   * having dependencies on himself.
   */
  @Test
  public void testActionScheduleWithDependencyOrderedCommandExecution() throws Exception {
    Type type = new TypeToken<Map<String, Set<String>>>() {}.getType();
    Map<String, List<String>> clusterHostInfo = StageUtils.getGson().fromJson(CLUSTER_HOST_INFO, type);

    Properties properties = new Properties();
    properties.setProperty("server.stage.command.execution_type", "DEPENDENCY_ORDERED");
    Configuration conf = new Configuration(properties);
    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);
    RoleCommandOrderProvider rcoProvider = mock(RoleCommandOrderProvider.class);
    RoleCommandOrder rco = mock(RoleCommandOrder.class);
    when(fsm.getCluster(anyString())).thenReturn(oneClusterMock);
    when(fsm.getClusterById(anyLong())).thenReturn(oneClusterMock);
    when(oneClusterMock.getService(anyString())).thenReturn(serviceObj);
    when(oneClusterMock.getClusterId()).thenReturn(Long.valueOf(1L));
    when(serviceObj.getServiceComponent(anyString())).thenReturn(scomp);
    when(scomp.getServiceComponentHost(anyString())).thenReturn(sch);
    when(serviceObj.getCluster()).thenReturn(oneClusterMock);

    when(rcoProvider.getRoleCommandOrder(1L)).thenReturn(rco);
    Map<RoleCommandPair, Set<RoleCommandPair>> roleCommandDependencies = new HashMap();
    RoleCommandPair roleCommand = new
            RoleCommandPair(Role.valueOf("NAMENODE"), RoleCommand.INSTALL);
    Set<RoleCommandPair> namenodeInstallDependencies = new HashSet<>();
    namenodeInstallDependencies.add(roleCommand);
    roleCommandDependencies.put(roleCommand, namenodeInstallDependencies);
    when(rco.getDependencies()).thenReturn(roleCommandDependencies);

    Host host = mock(Host.class);
    HashMap<String, ServiceComponentHost> hosts =
            new HashMap<>();
    hosts.put(hostname, sch);
    HostEntity hostEntity = new HostEntity();
    hostEntity.setHostName(hostname);
    hostEntity.setHostId(hostId);
    hostDAO.merge(hostEntity);
    when(scomp.getServiceComponentHosts()).thenReturn(hosts);

    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(ActionDBAccessorImpl.class);
    HostRoleCommandDAO hostRoleCommandDAOMock = mock(HostRoleCommandDAO.class);
    Mockito.doNothing().when(hostRoleCommandDAOMock).publishTaskCreateEvent(anyListOf(HostRoleCommand.class));
    Stage s = StageUtils.getATestStage(1, 977, hostname,
            "{\"host_param\":\"param_value\"}", "{\"stage_param\":\"param_value\"}");
    s.setCommandExecutionType(CommandExecutionType.DEPENDENCY_ORDERED);
    List<Stage> stages = Collections.singletonList(s);
    when(db.getCommandsInProgressCount()).thenReturn(stages.size());
    when(db.getFirstStageInProgressPerRequest()).thenReturn(stages);

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


    //Keep large number of attempts so that the task is not expired finally
    //Small action timeout to test rescheduling
    ActionScheduler scheduler = new ActionScheduler(100, 5, db, fsm,
            10000, new HostsMap((String) null), unitOfWork, null, conf,
            entityManagerProviderMock, hostRoleCommandDAOMock, null, rcoProvider, agentCommandsPublisher);
    scheduler.setTaskTimeoutAdjustment(false);

    List<AgentCommand> commands = waitForQueueSize(hostId, agentCommandsPublisher, 1, scheduler);

    assertTrue(commands != null && commands.size() == 1);
    AgentCommand scheduledCommand = commands.get(0);
    assertTrue(scheduledCommand instanceof ExecutionCommand);
    assertEquals("1-977", ((ExecutionCommand) scheduledCommand).getCommandId());
    assertEquals(clusterHostInfo, ((ExecutionCommand) scheduledCommand).getClusterHostInfo());

    //The action status has not changed, it should be queued again.
    commands = waitForQueueSize(hostId, agentCommandsPublisher, 2, scheduler);

    assertTrue(commands != null && commands.size() == 2);
    scheduledCommand = commands.get(1);
    assertTrue(scheduledCommand instanceof ExecutionCommand);
    assertEquals("1-977", ((ExecutionCommand) scheduledCommand).getCommandId());
    assertEquals(clusterHostInfo, ((ExecutionCommand) scheduledCommand).getClusterHostInfo());

    //Now change the action status
    s.setHostRoleStatus(hostname, "NAMENODE", HostRoleStatus.COMPLETED);

    //Wait for sometime, it shouldn't be scheduled this time.
    scheduler.doWork();

    EasyMock.verify(entityManagerProviderMock);
  }

  /**
   * Test whether scheduler times out an action
   */
  @Test
  public void testActionTimeout() throws Exception {
    Properties properties = new Properties();
    Configuration conf = new Configuration(properties);
    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);
    when(fsm.getCluster(nullable(String.class))).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);
    Host host = mock(Host.class);
    HashMap<String, ServiceComponentHost> hosts =
      new HashMap<>();
    hosts.put(hostname, sch);
    when(scomp.getServiceComponentHosts()).thenReturn(hosts);

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

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

    final Stage s = StageUtils.getATestStage(1, 977, hostname,
      "{\"host_param\":\"param_value\"}", "{\"stage_param\":\"param_value\"}");
    s.addHostRoleExecutionCommand(hostname, Role.SECONDARY_NAMENODE, RoleCommand.INSTALL,
            new ServiceComponentHostInstallEvent("SECONDARY_NAMENODE", hostname, System.currentTimeMillis(), "HDP-1.2.0"),
            "cluster1", "HDFS", false, false);
    s.setHostRoleStatus(hostname, "SECONDARY_NAMENODE", HostRoleStatus.IN_PROGRESS);
    List<Stage> stages = Collections.singletonList(s);

    ActionDBAccessor db = mock(ActionDBAccessor.class);
    HostRoleCommandDAO hostRoleCommandDAOMock = mock(HostRoleCommandDAO.class);
    Mockito.doNothing().when(hostRoleCommandDAOMock).publishTaskCreateEvent(anyListOf(HostRoleCommand.class));
    when(db.getCommandsInProgressCount()).thenReturn(stages.size());
    when(db.getFirstStageInProgressPerRequest()).thenReturn(stages);

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

    doAnswer(new Answer<Void>() {
      @Override
      public Void answer(InvocationOnMock invocation) throws Throwable {
        String host = (String) invocation.getArguments()[0];
        String role = (String) invocation.getArguments()[3];
        HostRoleCommand command = s.getHostRoleCommand(host, role);
        command.setStatus(HostRoleStatus.TIMEDOUT);
        return null;
      }
    }).when(db).timeoutHostRole(nullable(String.class), anyLong(), anyLong(), anyString(), anyBoolean(), eq(false));


    //Small action timeout to test rescheduling
    ActionScheduler scheduler = new ActionScheduler(100, 0, db, fsm, 3,
        new HostsMap((String) null), unitOfWork, null, conf, entityManagerProviderMock, hostRoleCommandDAOMock, null,
        agentCommandsPublisher);
    scheduler.setTaskTimeoutAdjustment(false);
    // Start the thread

    int cycleCount = 0;
    scheduler.doWork();
    //Check that in_progress command is rescheduled
    assertEquals(HostRoleStatus.QUEUED, stages.get(0).getHostRoleStatus(hostname, "SECONDARY_NAMENODE"));

    // Check was generated cancel command on timeout
    List<AgentCommand> commands = waitForQueueSize(hostId, agentCommandsPublisher, 1, scheduler);
    assertTrue(commands != null && commands.size() >= 1);
    AgentCommand scheduledCommand = commands.get(0);
    assertTrue(scheduledCommand instanceof CancelCommand);

    //Switch command back to IN_PROGRESS status and check that other command is not rescheduled
    stages.get(0).setHostRoleStatus(hostname, "SECONDARY_NAMENODE", HostRoleStatus.IN_PROGRESS);
    scheduler.doWork();
    assertEquals(2, stages.get(0).getAttemptCount(hostname, "NAMENODE"));
    assertEquals(3, stages.get(0).getAttemptCount(hostname, "SECONDARY_NAMENODE"));

    while (!stages.get(0).getHostRoleStatus(hostname, "SECONDARY_NAMENODE")
        .equals(HostRoleStatus.TIMEDOUT) && cycleCount++ <= MAX_CYCLE_ITERATIONS) {
      scheduler.doWork();
    }
    assertEquals(HostRoleStatus.TIMEDOUT,
            stages.get(0).getHostRoleStatus(hostname, "SECONDARY_NAMENODE"));

    verify(db, times(1)).startRequest(eq(1L));
    verify(db, times(1)).abortOperation(1L);

    EasyMock.verify(entityManagerProviderMock);
  }

  @Test
  public void testActionTimeoutForLostHost() throws Exception {
    Properties properties = new Properties();
    Configuration conf = new Configuration(properties);
    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);
    when(fsm.getCluster(nullable(String.class))).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);
    Host host = mock(Host.class);
    HashMap<String, ServiceComponentHost> hosts =
      new HashMap<>();
    hosts.put(hostname, sch);
    when(scomp.getServiceComponentHosts()).thenReturn(hosts);

    when(fsm.getHost(anyString())).thenReturn(host);
    when(host.getState()).thenReturn(HostState.HEARTBEAT_LOST);
    when(host.getHostName()).thenReturn(hostname);

    final Stage s = StageUtils.getATestStage(1, 977, hostname,
      "{\"host_param\":\"param_value\"}", "{\"stage_param\":\"param_value\"}");

    List<Stage> stages = Collections.singletonList(s);

    ActionDBAccessor db = mock(ActionDBAccessor.class);

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

    when(db.getCommandsInProgressCount()).thenReturn(stages.size());
    when(db.getFirstStageInProgressPerRequest()).thenReturn(stages);
    HostRoleCommandDAO hostRoleCommandDAOMock = mock(HostRoleCommandDAO.class);
    Mockito.doNothing().when(hostRoleCommandDAOMock).publishTaskCreateEvent(anyListOf(HostRoleCommand.class));

    doAnswer(new Answer<Void>() {
      @Override
      public Void answer(InvocationOnMock invocation) throws Throwable {
        String host = (String) invocation.getArguments()[0];
        String role = (String) invocation.getArguments()[3];
        HostRoleCommand command = s.getHostRoleCommand(host, role);
        command.setStatus(HostRoleStatus.ABORTED);
        return null;
      }
    }).when(db).timeoutHostRole(nullable(String.class), anyLong(), anyLong(), anyString(), anyBoolean(), eq(true));

    //Small action timeout to test rescheduling
    AmbariEventPublisher aep = EasyMock.createNiceMock(AmbariEventPublisher.class);
    ActionScheduler scheduler = new ActionScheduler(100, 0, db, fsm, 3,
        new HostsMap((String) null), unitOfWork, null, conf, entityManagerProviderMock, hostRoleCommandDAOMock, null,
        agentCommandsPublisher);
    scheduler.setTaskTimeoutAdjustment(false);

    int cycleCount=0;
    while (!stages.get(0).getHostRoleStatus(hostname, "NAMENODE")
      .equals(HostRoleStatus.ABORTED) && cycleCount++ <= MAX_CYCLE_ITERATIONS) {
      scheduler.doWork();
    }

    Assert.assertEquals(HostRoleStatus.ABORTED,stages.get(0).getHostRoleStatus(hostname, "NAMENODE"));

    EasyMock.verify(entityManagerProviderMock);
  }

  @Test
  public void testOpFailedEventRaisedForAbortedHostRole() throws Exception {
    Properties properties = new Properties();
    Configuration conf = new Configuration(properties);
    Clusters fsm = mock(Clusters.class);
    Cluster oneClusterMock = mock(Cluster.class);
    Service serviceObj = mock(Service.class);
    ServiceComponent scomp = mock(ServiceComponent.class);
    ServiceComponentHost sch1 = mock(ServiceComponentHost.class);
    ServiceComponentHost sch2 = mock(ServiceComponentHost.class);
    String hostname1 = "host1";
    String hostname2 = "host2";
    Host host1 = mock(Host.class);
    Host host2 = mock(Host.class);
    HostEntity hostEntity1 = new HostEntity();
    hostEntity1.setHostName(hostname1);
    HostEntity hostEntity2 = new HostEntity();
    hostEntity2.setHostName(hostname2);
    hostDAO.merge(hostEntity1);
    hostDAO.merge(hostEntity2);

    HashMap<String, ServiceComponentHost> hosts =
      new HashMap<>();
    hosts.put(hostname1, sch1);
    hosts.put(hostname2, sch2);
    when(scomp.getServiceComponentHosts()).thenReturn(hosts);

    when(fsm.getHost(hostname1)).thenReturn(host1);
    when(fsm.getHost(hostname2)).thenReturn(host2);
    when(host1.getState()).thenReturn(HostState.HEARTBEAT_LOST);
    when(host2.getState()).thenReturn(HostState.HEALTHY);
    when(host1.getHostName()).thenReturn(hostname1);
    when(host2.getHostName()).thenReturn(hostname2);
    when(scomp.getServiceComponentHost(hostname1)).thenReturn(sch1);
    when(scomp.getServiceComponentHost(hostname2)).thenReturn(sch2);

    UnitOfWork unitOfWork = mock(UnitOfWork.class);
    AgentCommandsPublisher agentCommandsPublisher = mock(AgentCommandsPublisher.class);
    when(fsm.getCluster(nullable(String.class))).thenReturn(oneClusterMock);
    when(oneClusterMock.getService(anyString())).thenReturn(serviceObj);
    when(serviceObj.getServiceComponent(anyString())).thenReturn(scomp);
    when(serviceObj.getCluster()).thenReturn(oneClusterMock);

    final Stage stage = stageFactory.createNew(1, "/tmp", "cluster1", 1L, "stageWith2Tasks",
      "{\"command_param\":\"param_value\"}", "{\"host_param\":\"param_value\"}");
    addInstallTaskToStage(stage, hostname1, "cluster1", Role.DATANODE,
      RoleCommand.INSTALL, Service.Type.HDFS, 1);
    addInstallTaskToStage(stage, hostname2, "cluster1", Role.NAMENODE,
      RoleCommand.INSTALL, Service.Type.HDFS, 2);
    final List<Stage> stages = Collections.singletonList(stage);

    ActionDBAccessor db = mock(ActionDBAccessor.class);

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

    when(db.getCommandsInProgressCount()).thenReturn(stages.size());
    when(db.getFirstStageInProgressPerRequest()).thenReturn(stages);
    HostRoleCommandDAO hostRoleCommandDAOMock = mock(HostRoleCommandDAO.class);
    Mockito.doNothing().when(hostRoleCommandDAOMock).publishTaskCreateEvent(anyListOf(HostRoleCommand.class));

    doAnswer(new Answer<Void>() {
      @Override
      public Void answer(InvocationOnMock invocation) throws Throwable {
        String host = (String) invocation.getArguments()[0];
        String role = (String) invocation.getArguments()[3];
        HostRoleCommand command = stage.getHostRoleCommand(host, role);
        command.setStatus(HostRoleStatus.ABORTED);
        return null;
      }
    }).when(db).timeoutHostRole(nullable(String.class), anyLong(), anyLong(), anyString(), anyBoolean(), eq(true));

    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());

    ArgumentCaptor<ServiceComponentHostEvent> eventsCapture1 =
      ArgumentCaptor.forClass(ServiceComponentHostEvent.class);
    ArgumentCaptor<ServiceComponentHostEvent> eventsCapture2 =
      ArgumentCaptor.forClass(ServiceComponentHostEvent.class);

    // Make sure the NN install doesn't timeout
    ActionScheduler scheduler = new ActionScheduler(100, 50000, db, fsm, 3,
        new HostsMap((String) null), unitOfWork, null, conf, entityManagerProviderMock,
        hostRoleCommandDAOMock, (HostRoleCommandFactory)null, agentCommandsPublisher);
