* 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.

*/

  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]}";

  /**

   * 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, List<String>> clusterHostInfo = StageUtils.getGson().fromJson(CLUSTER_HOST_INFO, type);

    ActionQueue aq = new ActionQueue();

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

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

    String hostname = "ahost.ambari.apache.org";

    Host host = mock(Host.class);

    HashMap<String, ServiceComponentHost> hosts =

            new HashMap<String, ServiceComponentHost>();

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

    ActionDBAccessor db = mock(ActionDBAccessorImpl.class);

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

    Stage s = StageUtils.getATestStage(1, 977, hostname, CLUSTER_HOST_INFO);

    stages.add(s);

    when(db.getStagesInProgress()).thenReturn(stages);

    //Keep large number of attempts so that the task is not expired finally

    //Small action timeout to test rescheduling

    ActionScheduler scheduler = new ActionScheduler(100, 100, db, aq, fsm,

        10000, new HostsMap((String) null), null, unitOfWork, conf);

    scheduler.setTaskTimeoutAdjustment(false);

    // Start the thread

    scheduler.start();

    List<AgentCommand> ac = waitForQueueSize(hostname, aq, 1);

    assertTrue(ac.get(0) instanceof ExecutionCommand);

    assertEquals("1-977", ((ExecutionCommand) (ac.get(0))).getCommandId());

    assertEquals(clusterHostInfo, ((ExecutionCommand) (ac.get(0))).getClusterHostInfo());

    //The action status has not changed, it should be queued again.

    ac = waitForQueueSize(hostname, aq, 1);

    assertTrue(ac.get(0) instanceof ExecutionCommand);

    assertEquals("1-977", ((ExecutionCommand) (ac.get(0))).getCommandId());

    assertEquals(clusterHostInfo, ((ExecutionCommand) (ac.get(0))).getClusterHostInfo());

    //Now change the action status

    s.setHostRoleStatus(hostname, "NAMENODE", HostRoleStatus.COMPLETED);

    ac = aq.dequeueAll(hostname);

    //Wait for sometime, it shouldn't be scheduled this time.

    ac = waitForQueueSize(hostname, aq, 0);

    scheduler.stop();

  }

  private List<AgentCommand> waitForQueueSize(String hostname, ActionQueue aq,

      int expectedQueueSize) throws InterruptedException {

    while (true) {

      List<AgentCommand> ac = aq.dequeueAll(hostname);     

      if (ac != null) {

        if (ac.size() == expectedQueueSize) {

          return ac;

        } else if (ac.size() > expectedQueueSize) {

          Assert.fail("Expected size : " + expectedQueueSize + " Actual size="

              + ac.size());

        }

      }

      Thread.sleep(100);

    }

  }

  /**

   * Test whether scheduler times out an action

   */

  @Test

  public void testActionTimeout() throws Exception {

    ActionQueue aq = new ActionQueue();

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

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

    String hostname = "ahost.ambari.apache.org";

    Host host = mock(Host.class);

    HashMap<String, ServiceComponentHost> hosts =

            new HashMap<String, ServiceComponentHost>();

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

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

    final Stage s = StageUtils.getATestStage(1, 977, hostname, CLUSTER_HOST_INFO);

    stages.add(s);

    ActionDBAccessor db = mock(ActionDBAccessor.class);

    when(db.getStagesInProgress()).thenReturn(stages);

    doAnswer(new Answer() {

      @Override

      public Object 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(anyString(), anyLong(), anyLong(), anyString());

    //Small action timeout to test rescheduling

    ActionScheduler scheduler = new ActionScheduler(100, 50, db, aq, fsm, 3,

        new HostsMap((String) null), null, unitOfWork, conf);

    scheduler.setTaskTimeoutAdjustment(false);

    // Start the thread

    scheduler.start();

    while (!stages.get(0).getHostRoleStatus(hostname, "NAMENODE")

        .equals(HostRoleStatus.TIMEDOUT)) {

      Thread.sleep(100);

    }

    assertEquals(stages.get(0).getHostRoleStatus(hostname, "NAMENODE"),

        HostRoleStatus.TIMEDOUT);

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

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

    scheduler.stop();

  }

  @Test

  public void testActionTimeoutForLostHost() throws Exception {

    ActionQueue aq = new ActionQueue();

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

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

    String hostname = "ahost.ambari.apache.org";

    Host host = mock(Host.class);

    HashMap<String, ServiceComponentHost> hosts =

            new HashMap<String, ServiceComponentHost>();

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

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

    final Stage s = StageUtils.getATestStage(1, 977, hostname, CLUSTER_HOST_INFO);

    stages.add(s);

    ActionDBAccessor db = mock(ActionDBAccessor.class);

    when(db.getStagesInProgress()).thenReturn(stages);

    doAnswer(new Answer() {

      @Override

      public Object 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(anyString(), anyLong(), anyLong(), anyString());

    //Small action timeout to test rescheduling

    ActionScheduler scheduler = new ActionScheduler(100, 50, db, aq, fsm, 3,

      new HostsMap((String) null), null, unitOfWork, conf);

    scheduler.setTaskTimeoutAdjustment(false);

    // Start the thread

    scheduler.start();

    while (!stages.get(0).getHostRoleStatus(hostname, "NAMENODE")

      .equals(HostRoleStatus.TIMEDOUT)) {

      Thread.sleep(100L);

    }

//      HostRoleStatus.TIMEDOUT);

    scheduler.stop();

  }

  @Test

  public void testOpFailedEventRaisedForAbortedHostRole() throws Exception {

    ActionQueue aq = new ActionQueue();

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

    HashMap<String, ServiceComponentHost> hosts =

            new HashMap<String, ServiceComponentHost>();

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

    when(fsm.getCluster(anyString())).thenReturn(oneClusterMock);

    when(oneClusterMock.getService(anyString())).thenReturn(serviceObj);

    when(serviceObj.getServiceComponent(anyString())).thenReturn(scomp);

    when(serviceObj.getCluster()).thenReturn(oneClusterMock);

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

    Stage stage = new Stage(1, "/tmp", "cluster1", "stageWith2Tasks",

      CLUSTER_HOST_INFO);

    addInstallTaskToStage(stage, hostname1, "cluster1", Role.DATANODE,

      RoleCommand.INSTALL, Service.Type.HDFS, 1);

    addInstallTaskToStage(stage, hostname2, "cluster1", Role.NAMENODE,

      RoleCommand.INSTALL, Service.Type.HDFS, 2);

    stages.add(stage);

    ActionDBAccessor db = mock(ActionDBAccessor.class);

    when(db.getStagesInProgress()).thenReturn(stages);

    doAnswer(new Answer() {

      @Override

      public Object answer(InvocationOnMock invocation) throws Throwable {

        String host = (String) invocation.getArguments()[0];

        String role = (String) invocation.getArguments()[3];

        //HostRoleCommand command = stages.get(0).getHostRoleCommand(host, role);

        for (HostRoleCommand command : stages.get(0).getOrderedHostRoleCommands()) {

          if (command.getHostName().equals(host) && command.getRole().name()

              .equals(role)) {

            command.setStatus(HostRoleStatus.TIMEDOUT);

          }

        }

        return null;

      }

    }).when(db).timeoutHostRole(anyString(), anyLong(), anyLong(), anyString());

    doAnswer(new Answer() {

      @Override

      public Object answer(InvocationOnMock invocation) throws Throwable {

        Long requestId = (Long) invocation.getArguments()[0];

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

              }

            }

          }

        }

        return null;

      }

    }).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, aq, fsm, 3,

      new HostsMap((String) null), null, unitOfWork, conf);

    scheduler.setTaskTimeoutAdjustment(false);

    // Start the thread

    scheduler.start();

    while (!stages.get(0).getHostRoleStatus(hostname1, "DATANODE")

      .equals(HostRoleStatus.TIMEDOUT) && !stages.get(0).getHostRoleStatus

      (hostname2, "NAMENODE").equals(HostRoleStatus.ABORTED)) {

      Thread.sleep(100L);

    }

    Assert.assertEquals(HostRoleStatus.TIMEDOUT,

      stages.get(0).getHostRoleStatus(hostname1, "DATANODE"));

    Assert.assertEquals(HostRoleStatus.ABORTED,

      stages.get(0).getHostRoleStatus(hostname2, "NAMENODE"));

    verify(sch1, atLeastOnce()).handleEvent(eventsCapture1.capture());

    verify(sch2, atLeastOnce()).handleEvent(eventsCapture2.capture());

    List<ServiceComponentHostEvent> eventTypes = eventsCapture1.getAllValues();

    eventTypes.addAll(eventsCapture2.getAllValues());

    Assert.assertNotNull(eventTypes);

    ServiceComponentHostOpFailedEvent datanodeFailedEvent = null;

    ServiceComponentHostOpFailedEvent namenodeFailedEvent = null;

    for (ServiceComponentHostEvent eventType : eventTypes) {

      if (eventType instanceof ServiceComponentHostOpFailedEvent) {

        ServiceComponentHostOpFailedEvent event =

          (ServiceComponentHostOpFailedEvent) eventType;

        if (event.getServiceComponentName().equals("DATANODE")) {

          datanodeFailedEvent = event;

        } else if (event.getServiceComponentName().equals("NAMENODE")) {

          namenodeFailedEvent = event;

        }

      }

    }

    Assert.assertNotNull("Datanode should be in Install failed state.",

      datanodeFailedEvent);

    Assert.assertNotNull("Namenode should be in Install failed state.",

      namenodeFailedEvent);

    scheduler.stop();

  }

  /**

   * Test server action

   */

  @Test

  public void testServerAction() throws Exception {

    ActionQueue aq = new ActionQueue();

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

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

    String hostname = "ahost.ambari.apache.org";

    HashMap<String, ServiceComponentHost> hosts =

            new HashMap<String, ServiceComponentHost>();

    hosts.put(hostname, sch);

    when(scomp.getServiceComponentHosts()).thenReturn(hosts);

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

    Map<String, String> payload = new HashMap<String, String>();

    payload.put(ServerAction.PayloadName.CLUSTER_NAME, "cluster1");

    payload.put(ServerAction.PayloadName.CURRENT_STACK_VERSION, "HDP-0.2");

    final Stage s = getStageWithServerAction(1, 977, hostname, payload, "test");

    stages.add(s);

    ActionDBAccessor db = mock(ActionDBAccessor.class);

    when(db.getStagesInProgress()).thenReturn(stages);

    doAnswer(new Answer() {

      @Override

      public Object answer(InvocationOnMock invocation) throws Throwable {

        String host = (String) invocation.getArguments()[0];

        String role = (String) invocation.getArguments()[3];

        CommandReport commandReport = (CommandReport) invocation.getArguments()[4];

        HostRoleCommand command = s.getHostRoleCommand(host, role);

        command.setStatus(HostRoleStatus.valueOf(commandReport.getStatus()));

        return null;

      }

    }).when(db).updateHostRoleState(anyString(), anyLong(), anyLong(), anyString(), any(CommandReport.class));

    ActionScheduler scheduler = new ActionScheduler(100, 50, db, aq, fsm, 3,

        new HostsMap((String) null), new ServerActionManagerImpl(fsm),

        unitOfWork, conf);

    scheduler.start();

    while (!stages.get(0).getHostRoleStatus(hostname, "AMBARI_SERVER_ACTION")

        .equals(HostRoleStatus.COMPLETED)) {

      Thread.sleep(100);

    }

    scheduler.stop();

    assertEquals(stages.get(0).getHostRoleStatus(hostname, "AMBARI_SERVER_ACTION"),

        HostRoleStatus.COMPLETED);

  }

  @Test

  public void testServerActionFailed() throws Exception {

    ActionQueue aq = new ActionQueue();

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

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

    String hostname = "ahost.ambari.apache.org";

    HashMap<String, ServiceComponentHost> hosts =

            new HashMap<String, ServiceComponentHost>();

    hosts.put(hostname, sch);

    when(scomp.getServiceComponentHosts()).thenReturn(hosts);

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

    Map<String, String> payload = new HashMap<String, String>();

    payload.put(ServerAction.PayloadName.CURRENT_STACK_VERSION, "HDP-0.2");

    final Stage s = getStageWithServerAction(1, 977, hostname, payload, "test");

    stages.add(s);

    ActionDBAccessor db = mock(ActionDBAccessor.class);

    when(db.getStagesInProgress()).thenReturn(stages);

    doAnswer(new Answer() {

      @Override

      public Object answer(InvocationOnMock invocation) throws Throwable {

        String host = (String) invocation.getArguments()[0];

        String role = (String) invocation.getArguments()[3];

        CommandReport commandReport = (CommandReport) invocation.getArguments()[4];

        HostRoleCommand command = s.getHostRoleCommand(host, role);

        command.setStatus(HostRoleStatus.valueOf(commandReport.getStatus()));

        return null;

      }

    }).when(db).updateHostRoleState(anyString(), anyLong(), anyLong(), anyString(), any(CommandReport.class));

    ActionScheduler scheduler = new ActionScheduler(100, 50, db, aq, fsm, 3,

        new HostsMap((String) null), new ServerActionManagerImpl(fsm), unitOfWork, conf);

    scheduler.start();

    while (!stages.get(0).getHostRoleStatus(hostname, "AMBARI_SERVER_ACTION")

        .equals(HostRoleStatus.FAILED)) {

      Thread.sleep(100);

    }

    scheduler.stop();

    assertEquals(stages.get(0).getHostRoleStatus(hostname, "AMBARI_SERVER_ACTION"),

        HostRoleStatus.FAILED);

    assertEquals("test", stages.get(0).getRequestContext());

  }

  private static Stage getStageWithServerAction(long requestId, long stageId, String hostName,

                                                Map<String, String> payload, String requestContext) {

    Stage stage = new Stage(requestId, "/tmp", "cluster1", requestContext, CLUSTER_HOST_INFO);

    stage.setStageId(stageId);

    long now = System.currentTimeMillis();

    stage.addServerActionCommand(ServerAction.Command.FINALIZE_UPGRADE, Role.AMBARI_SERVER_ACTION,

        RoleCommand.EXECUTE, "cluster1",

        new ServiceComponentHostUpgradeEvent("AMBARI_SERVER_ACTION", hostName, now, "HDP-0.2"),

        hostName);

    ExecutionCommand execCmd = stage.getExecutionCommandWrapper(hostName,

        Role.AMBARI_SERVER_ACTION.toString()).getExecutionCommand();

    execCmd.setCommandParams(payload);

    return stage;

  }

  /**

   * Verifies that stages that are executed on different hosts and

   * rely to different requests are scheduled to be  executed in parallel

   */

  @Test

  public void testIndependentStagesExecution() throws Exception {

    ActionQueue aq = new ActionQueue();

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

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

    String hostname1 = "ahost.ambari.apache.org";

    String hostname2 = "bhost.ambari.apache.org";

    String hostname3 = "chost.ambari.apache.org";

    String hostname4 = "chost.ambari.apache.org";

    HashMap<String, ServiceComponentHost> hosts =

            new HashMap<String, ServiceComponentHost>();

    hosts.put(hostname1, sch);

    hosts.put(hostname2, sch);

    hosts.put(hostname3, sch);

    hosts.put(hostname4, sch);

    when(scomp.getServiceComponentHosts()).thenReturn(hosts);

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

    stages.add(

            getStageWithSingleTask(

                    hostname1, "cluster1", Role.DATANODE,

                    RoleCommand.START, Service.Type.HDFS, 1, 1, 1));

    stages.add( // Stage with the same hostname, should not be scheduled

            getStageWithSingleTask(

                    hostname1, "cluster1", Role.GANGLIA_MONITOR,

                    RoleCommand.START, Service.Type.GANGLIA, 2, 2, 2));

    stages.add(

            getStageWithSingleTask(

                    hostname2, "cluster1", Role.DATANODE,

                    RoleCommand.START, Service.Type.HDFS, 3, 3, 3));

    stages.add(

            getStageWithSingleTask(

                    hostname3, "cluster1", Role.DATANODE,

                    RoleCommand.START, Service.Type.HDFS, 4, 4, 4));

    stages.add( // Stage with the same request id, should not be scheduled

            getStageWithSingleTask(

                    hostname4, "cluster1", Role.GANGLIA_MONITOR,

                    RoleCommand.START, Service.Type.GANGLIA, 5, 5, 4));

    ActionDBAccessor db = mock(ActionDBAccessor.class);

    when(db.getStagesInProgress()).thenReturn(stages);

    Properties properties = new Properties();

    Configuration conf = new Configuration(properties);

    ActionScheduler scheduler = new ActionScheduler(100, 50, db, aq, fsm, 3,

            new HostsMap((String) null), new ServerActionManagerImpl(fsm),

            unitOfWork, conf);

    ActionManager am = new ActionManager(

            2, 2, aq, fsm, db, new HostsMap((String) null),

            new ServerActionManagerImpl(fsm), unitOfWork, requestFactory, conf);

    scheduler.doWork();

    Assert.assertEquals(HostRoleStatus.QUEUED, stages.get(0).getHostRoleStatus(hostname1, "DATANODE"));

    Assert.assertEquals(HostRoleStatus.PENDING, stages.get(1).getHostRoleStatus(hostname1, "GANGLIA_MONITOR"));

    Assert.assertEquals(HostRoleStatus.QUEUED, stages.get(2).getHostRoleStatus(hostname2, "DATANODE"));

    Assert.assertEquals(HostRoleStatus.QUEUED, stages.get(3).getHostRoleStatus(hostname3, "DATANODE"));

    Assert.assertEquals(HostRoleStatus.PENDING, stages.get(4).getHostRoleStatus(hostname4, "GANGLIA_MONITOR"));

  }

  /**

   * Verifies that ActionScheduler respects "disable parallel stage execution option"

   */

  @Test

  public void testIndependentStagesExecutionDisabled() throws Exception {

    ActionQueue aq = new ActionQueue();

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

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

    String hostname1 = "ahost.ambari.apache.org";

    String hostname2 = "bhost.ambari.apache.org";

    String hostname3 = "chost.ambari.apache.org";

    String hostname4 = "chost.ambari.apache.org";

    HashMap<String, ServiceComponentHost> hosts =

            new HashMap<String, ServiceComponentHost>();

    hosts.put(hostname1, sch);

    hosts.put(hostname2, sch);

    hosts.put(hostname3, sch);

    hosts.put(hostname4, sch);

    when(scomp.getServiceComponentHosts()).thenReturn(hosts);

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

    stages.add(

            getStageWithSingleTask(

                    hostname1, "cluster1", Role.DATANODE,

                    RoleCommand.START, Service.Type.HDFS, 1, 1, 1));

    stages.add( // Stage with the same hostname, should not be scheduled

            getStageWithSingleTask(

                    hostname1, "cluster1", Role.GANGLIA_MONITOR,

                    RoleCommand.START, Service.Type.GANGLIA, 2, 2, 2));

    stages.add(

            getStageWithSingleTask(

                    hostname2, "cluster1", Role.DATANODE,

                    RoleCommand.START, Service.Type.HDFS, 3, 3, 3));

    stages.add(

            getStageWithSingleTask(

                    hostname3, "cluster1", Role.DATANODE,

                    RoleCommand.START, Service.Type.HDFS, 4, 4, 4));

    stages.add( // Stage with the same request id, should not be scheduled

            getStageWithSingleTask(

                    hostname4, "cluster1", Role.GANGLIA_MONITOR,

                    RoleCommand.START, Service.Type.GANGLIA, 5, 5, 4));

    ActionDBAccessor db = mock(ActionDBAccessor.class);

    when(db.getStagesInProgress()).thenReturn(stages);

    Properties properties = new Properties();

    properties.put(Configuration.PARALLEL_STAGE_EXECUTION_KEY, "false");

    Configuration conf = new Configuration(properties);

    ActionScheduler scheduler = new ActionScheduler(100, 50, db, aq, fsm, 3,

            new HostsMap((String) null), new ServerActionManagerImpl(fsm),

            unitOfWork, conf);

    ActionManager am = new ActionManager(

            2, 2, aq, fsm, db, new HostsMap((String) null),

            new ServerActionManagerImpl(fsm), unitOfWork,

            requestFactory, conf);

    scheduler.doWork();

    Assert.assertEquals(HostRoleStatus.QUEUED, stages.get(0).getHostRoleStatus(hostname1, "DATANODE"));

    Assert.assertEquals(HostRoleStatus.PENDING, stages.get(1).getHostRoleStatus(hostname1, "GANGLIA_MONITOR"));

    Assert.assertEquals(HostRoleStatus.PENDING, stages.get(2).getHostRoleStatus(hostname2, "DATANODE"));

    Assert.assertEquals(HostRoleStatus.PENDING, stages.get(3).getHostRoleStatus(hostname3, "DATANODE"));

    Assert.assertEquals(HostRoleStatus.PENDING, stages.get(4).getHostRoleStatus(hostname4, "GANGLIA_MONITOR"));

  }

  @Test

  public void testRequestFailureOnStageFailure() throws Exception {

    ActionQueue aq = new ActionQueue();

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

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

    String hostname = "ahost.ambari.apache.org";

    HashMap<String, ServiceComponentHost> hosts =

            new HashMap<String, ServiceComponentHost>();

    hosts.put(hostname, sch);

    when(scomp.getServiceComponentHosts()).thenReturn(hosts);

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

    stages.add(

        getStageWithSingleTask(

            hostname, "cluster1", Role.NAMENODE, RoleCommand.UPGRADE, Service.Type.HDFS, 1, 1, 1));

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

    when(db.getStagesInProgress()).thenReturn(stages);

    doAnswer(new Answer() {

      @Override

      public Object 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];

          String role = report.getRole();

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

      @Override

      public Object answer(InvocationOnMock invocation) throws Throwable {

        Long requestId = (Long) invocation.getArguments()[0];

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

              }

            }

          }

        }

        return null;

      }

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

    Properties properties = new Properties();

    Configuration conf = new Configuration(properties);

    ActionScheduler scheduler = new ActionScheduler(100, 50, db, aq, fsm, 3,

        new HostsMap((String) null), new ServerActionManagerImpl(fsm),

        unitOfWork, conf);

    ActionManager am = new ActionManager(

        2, 2, aq, fsm, db, new HostsMap((String) null), new ServerActionManagerImpl(fsm), unitOfWork, requestFactory, conf);

    scheduler.doWork();

    List<CommandReport> reports = new ArrayList<CommandReport>();

    reports.add(getCommandReport(HostRoleStatus.FAILED, Role.NAMENODE, Service.Type.HDFS, "1-1", 1));

    am.processTaskResponse(hostname, reports);

    scheduler.doWork();

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

    Assert.assertEquals(HostRoleStatus.ABORTED, stages.get(1).getHostRoleStatus(hostname, "DATANODE"));

  }

  /**

   * Tests that the whole request is aborted when there are no QUEUED tasks for a role and

   * success factor is not met. As long as there is one QUEUED task the request is not

   * aborted.

   * @throws Exception

   */

  @Test

  public void testRequestAbortsOnlyWhenNoQueuedTaskAndSuccessFactorUnmet() throws Exception {

    ActionQueue aq = new ActionQueue();

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

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

    String host1 = "host1";

    String host2 = "host2";

    Host host = mock(Host.class);

    HashMap<String, ServiceComponentHost> hosts =

            new HashMap<String, ServiceComponentHost>();

    hosts.put(host1, sch);

    hosts.put(host2, sch);

    when(scomp.getServiceComponentHosts()).thenReturn(hosts);

    when(fsm.getHost(anyString())).thenReturn(host);

    when(host.getState()).thenReturn(HostState.HEALTHY);

    when(host.getHostName()).thenReturn(host1);

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

    long now = System.currentTimeMillis();

    Stage stage = new Stage(1, "/tmp", "cluster1", "testRequestFailureBasedOnSuccessFactor", CLUSTER_HOST_INFO);

    stage.setStageId(1);

    addHostRoleExecutionCommand(now, stage, Role.SQOOP, Service.Type.SQOOP,

        RoleCommand.INSTALL, host1, "cluster1");

    addHostRoleExecutionCommand(now, stage, Role.OOZIE_CLIENT, Service.Type.OOZIE,

        RoleCommand.INSTALL, host1, "cluster1");

    addHostRoleExecutionCommand(now, stage, Role.MAPREDUCE_CLIENT, Service.Type.MAPREDUCE,

        RoleCommand.INSTALL, host1, "cluster1");

    addHostRoleExecutionCommand(now, stage, Role.HBASE_CLIENT, Service.Type.HBASE,

        RoleCommand.INSTALL, host1, "cluster1");

    addHostRoleExecutionCommand(now, stage, Role.GANGLIA_MONITOR, Service.Type.GANGLIA,

        RoleCommand.INSTALL, host1, "cluster1");

    addHostRoleExecutionCommand(now, stage, Role.HBASE_CLIENT, Service.Type.HBASE,

        RoleCommand.INSTALL, host2, "cluster1");

    addHostRoleExecutionCommand(now, stage, Role.GANGLIA_MONITOR, Service.Type.GANGLIA,

        RoleCommand.INSTALL, host2, "cluster1");

    stages.add(stage);

    HostRoleStatus[] statusesAtIterOne = {HostRoleStatus.QUEUED, HostRoleStatus.QUEUED,

        HostRoleStatus.QUEUED, HostRoleStatus.QUEUED, HostRoleStatus.FAILED,

        HostRoleStatus.FAILED, HostRoleStatus.QUEUED, HostRoleStatus.QUEUED};

    for (int index = 0; index < stage.getOrderedHostRoleCommands().size(); index++) {

      stage.getOrderedHostRoleCommands().get(index).setTaskId(index + 1);

      stage.getOrderedHostRoleCommands().get(index).setStatus(statusesAtIterOne[index]);

    }

    stage.setLastAttemptTime(host1, Role.SQOOP.toString(), now);

    stage.setLastAttemptTime(host1, Role.MAPREDUCE_CLIENT.toString(), now);

    stage.setLastAttemptTime(host1, Role.OOZIE_CLIENT.toString(), now);

    stage.setLastAttemptTime(host1, Role.GANGLIA_MONITOR.toString(), now);

    stage.setLastAttemptTime(host1, Role.HBASE_CLIENT.toString(), now);

    stage.setLastAttemptTime(host2, Role.GANGLIA_MONITOR.toString(), now);

    stage.setLastAttemptTime(host2, Role.HBASE_CLIENT.toString(), now);

    ActionDBAccessor db = mock(ActionDBAccessor.class);

    when(db.getStagesInProgress()).thenReturn(stages);

    doAnswer(new Answer() {

      @Override

      public Object answer(InvocationOnMock invocation) throws Throwable {

        String host = (String) invocation.getArguments()[0];

        Long requestId = (Long) invocation.getArguments()[1];

        Long stageId = (Long) invocation.getArguments()[2];

        String role = (String) invocation.getArguments()[3];

        CommandReport commandReport = (CommandReport) invocation.getArguments()[4];

        for (Stage stage : stages) {

          if (requestId.equals(stage.getRequestId()) && stageId.equals(stage.getStageId())) {

            HostRoleCommand command = stage.getHostRoleCommand(host, role);

            command.setStatus(HostRoleStatus.valueOf(commandReport.getStatus()));

          }

        }

        return null;

      }

    }).when(db).updateHostRoleState(anyString(), anyLong(), anyLong(), anyString(), any(CommandReport.class));

    when(db.getTask(anyLong())).thenAnswer(new Answer<Object>() {

      @Override

      public Object answer(InvocationOnMock invocation) throws Throwable {