AMBARI-11891. Add Service fails after successful RU due to yum cache (ncole)

ambari-server/src/main/java/org/apache/ambari/server/controller/internal/UpgradeResourceProvider.java

@@ -399,7 +399,7 @@ public class UpgradeResourceProvider extends AbstractControllerResourceProvider
       List<Long> taskIds = new ArrayList<Long>();
 
       for (HostRoleCommand hrc : internalRequest.getCommands()) {
-        if (HostRoleStatus.ABORTED == hrc.getStatus()) {
+        if (HostRoleStatus.ABORTED == hrc.getStatus() || HostRoleStatus.TIMEDOUT == hrc.getStatus()) {
           taskIds.add(hrc.getTaskId());
         }
       }
@@ -556,7 +556,7 @@ public class UpgradeResourceProvider extends AbstractControllerResourceProvider
     StackId sourceStackId = null;
     StackId targetStackId = null;
 
-    switch( direction ){
+    switch (direction) {
       case UPGRADE:
         sourceStackId = cluster.getCurrentStackVersion();
 
@@ -850,12 +850,15 @@ public class UpgradeResourceProvider extends AbstractControllerResourceProvider
     // service_package_folder and hooks_folder params.
     AmbariMetaInfo ambariMetaInfo = s_metaProvider.get();
     StackId stackId = cluster.getDesiredStackVersion();
+
+
     StackInfo stackInfo = ambariMetaInfo.getStack(stackId.getStackName(), stackId.getStackVersion());
-    if (wrapper.getTasks() != null && wrapper.getTasks().size() > 0) {
+    if (wrapper.getTasks() != null &&
+        wrapper.getTasks().size() > 0 &&
+        wrapper.getTasks().get(0).getService() != null) {
       String serviceName = wrapper.getTasks().get(0).getService();
       ServiceInfo serviceInfo = ambariMetaInfo.getService(stackId.getStackName(), stackId.getStackVersion(), serviceName);
-      params.put(SERVICE_PACKAGE_FOLDER,
-          serviceInfo.getServicePackageFolder());
+      params.put(SERVICE_PACKAGE_FOLDER, serviceInfo.getServicePackageFolder());
       params.put(HOOKS_FOLDER, stackInfo.getStackHooksFolder());
     }
 

ambari-server/src/main/java/org/apache/ambari/server/state/UpgradeHelper.java

@@ -305,8 +305,17 @@ public class UpgradeHelper {
     holder.title = tokenReplace(ctx, holder.title, null, null);
 
     for (StageWrapper stageWrapper : holder.items) {
+      if (null != stageWrapper.getText()) {
+        stageWrapper.setText(tokenReplace(ctx, stageWrapper.getText(),
+            null, null));
+      }
+
       for (TaskWrapper taskWrapper : stageWrapper.getTasks()) {
         for (Task task : taskWrapper.getTasks()) {
+          if (null != task.summary) {
+            task.summary = tokenReplace(ctx, task.summary, null, null);
+          }
+
           if (task.getType() == Type.MANUAL) {
             ManualTask mt = (ManualTask) task;
             if (null != mt.message) {

ambari-server/src/main/java/org/apache/ambari/server/state/stack/upgrade/ClusterGrouping.java

@@ -19,6 +19,7 @@ package org.apache.ambari.server.state.stack.upgrade;
 
 import java.util.ArrayList;
 import java.util.Collections;
+import java.util.HashSet;
 import java.util.LinkedHashSet;
 import java.util.List;
 import java.util.Map;
@@ -34,6 +35,7 @@ import javax.xml.bind.annotation.XmlTransient;
 import javax.xml.bind.annotation.XmlType;
 
 import org.apache.ambari.server.stack.HostsType;
+import org.apache.ambari.server.state.Host;
 import org.apache.ambari.server.state.UpgradeContext;
 import org.apache.ambari.server.state.stack.UpgradePack.ProcessingComponent;
 
@@ -192,7 +194,7 @@ public class ClusterGrouping extends Grouping {
   private StageWrapper getExecuteStageWrapper(UpgradeContext ctx, ExecuteStage execution) {
     String service   = execution.service;
     String component = execution.component;
-    Task task        = execution.task;
+    ExecuteTask et = (ExecuteTask) execution.task;
 
     if (null != service && !service.isEmpty() &&
         null != component && !component.isEmpty()) {
@@ -202,8 +204,6 @@ public class ClusterGrouping extends Grouping {
       if (hosts != null) {
         Set<String> realHosts = new LinkedHashSet<String>(hosts.hosts);
 
-        ExecuteTask et = (ExecuteTask) task;
-
         if (null != et.hosts && "master".equals(et.hosts) && null != hosts.master) {
           realHosts = Collections.singleton(hosts.master);
         }
@@ -211,8 +211,20 @@ public class ClusterGrouping extends Grouping {
         return new StageWrapper(
             StageWrapper.Type.RU_TASKS,
             execution.title,
-            new TaskWrapper(service, component, realHosts, task));
+            new TaskWrapper(service, component, realHosts, et));
+      }
+    } else if (null == service && null == component) {
+      // no service, no component goes to all hosts
+
+      Set<String> hostNames = new HashSet<String>();
+      for (Host host : ctx.getCluster().getHosts()) {
+        hostNames.add(host.getHostName());
       }
+
+      return new StageWrapper(
+          StageWrapper.Type.RU_TASKS, execution.title,
+          new TaskWrapper(service, component, hostNames, et));
+
     }
     return null;
   }

ambari-server/src/main/java/org/apache/ambari/server/state/stack/upgrade/StageWrapper.java

@@ -91,6 +91,13 @@ public class StageWrapper {
     return text;
   }
 
+  /**
+   * @param text the new text for the stage
+   */
+  public void setText(String newText) {
+    text = newText;
+  }
+
   /**
    * Gets the type of stage.  All tasks defined for the stage execute this type.
    * @return the type

ambari-server/src/main/resources/custom_actions/scripts/ru_execute_tasks.py

@@ -110,15 +110,23 @@ class ExecuteUpgradeTasks(Script):
         Logger.info(str(task))
 
         # If a (script, function) exists, it overwrites the command.
-        if task.script and task.function and service_package_folder and hooks_folder:
+        if task.script and task.function:
           file_cache = FileCache(agent_config)
-          command_paths = {"commandParams":
-                                 {"service_package_folder": service_package_folder,
-                                  "hooks_folder": hooks_folder
-                                 }
-                              }
+
           server_url_prefix = default('/hostLevelParams/jdk_location', "")
-          base_dir = file_cache.get_service_base_dir(command_paths, server_url_prefix)
+
+          if service_package_folder and hooks_folder:
+            command_paths = {
+              "commandParams": {
+                "service_package_folder": service_package_folder,
+                "hooks_folder": hooks_folder
+              }
+            } 
+
+            base_dir = file_cache.get_service_base_dir(command_paths, server_url_prefix)
+          else:
+            base_dir = file_cache.get_custom_actions_base_dir(server_url_prefix)
+
           script_path = os.path.join(base_dir, task.script)
           if not os.path.exists(script_path):
             message = "Script %s does not exist" % str(script_path)

ambari-server/src/main/resources/custom_actions/scripts/ru_set_all.py

@@ -0,0 +1,62 @@
+#!/usr/bin/env python
+"""
+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.
+
+Ambari Agent
+
+"""
+
+from ambari_commons.os_check import OSCheck
+from resource_management.libraries.script import Script
+from resource_management.libraries.functions.default import default
+from resource_management.libraries.functions.version import compare_versions
+from resource_management.libraries.functions.version import format_hdp_stack_version
+from resource_management.core import shell
+from resource_management.core.exceptions import Fail
+from resource_management.core.logger import Logger
+
+class UpgradeSetAll(Script):
+  """
+  This script is a part of Rolling Upgrade workflow and is used to set the
+  component versions as a final step in the upgrade process
+  """
+
+  def actionexecute(self, env):
+    config = Script.get_config()
+
+    version = default('/commandParams/version', None)
+    stack_name = default('/hostLevelParams/stack_name', "")
+
+    if not version:
+      raise Fail("Value is required for '/commandParams/version'")
+  
+    # other os?
+    if OSCheck.is_redhat_family():
+      cmd = "/usr/bin/yum clean all"
+      code, out = shell.call(cmd)
+      Logger.info("Command: {0}\nCode: {1}, Out: {2}".format(cmd, str(code), str(out)))
+
+    min_ver = format_hdp_stack_version("2.2")
+    real_ver = format_hdp_stack_version(version)
+    if stack_name == "HDP":
+      if compare_versions(real_ver, min_ver) >= 0:
+        cmd = "hdp-select set all {0}".format(version)
+        code, out = shell.call(cmd)
+        Logger.info("Command: {0}\nCode: {1}, Out: {2}".format(cmd, str(code), str(out)))
+
+if __name__ == "__main__":
+  UpgradeSetAll().execute()

ambari-server/src/main/resources/stacks/HDP/2.2/upgrades/upgrade-2.2.xml

@@ -288,6 +288,14 @@
           <function>finalize_rolling_upgrade</function>
         </task>
       </execute-stage>
+      
+      <execute-stage title="Update remaining HDP stack to {{version}}">
+        <task xsi:type="execute">
+          <script>scripts/ru_set_all.py</script>
+          <function>actionexecute</function>
+        </task>
+      </execute-stage>
+      
       <execute-stage title="Save Cluster State" service="" component="">
         <task xsi:type="server_action" class="org.apache.ambari.server.serveraction.upgrades.FinalizeUpgradeAction">
         </task>

ambari-server/src/main/resources/stacks/HDP/2.2/upgrades/upgrade-2.3.xml

@@ -23,6 +23,7 @@
   <order>
     <group xsi:type="cluster" name="PRE_CLUSTER" title="Prepare Backups">
       <direction>UPGRADE</direction>
+      
       <execute-stage service="HDFS" component="NAMENODE" title="Pre Upgrade HDFS">
         <task xsi:type="execute" hosts="master">
           <script>scripts/namenode.py</script>
@@ -299,10 +300,17 @@
         </task>
       </execute-stage>
 
-      <execute-stage title="Save Cluster State" service="" component="">
-        <task xsi:type="server_action" class="org.apache.ambari.server.serveraction.upgrades.FinalizeUpgradeAction">
+      <execute-stage title="Update remaining HDP stack to {{version}}">
+        <task xsi:type="execute">
+          <script>scripts/ru_set_all.py</script>
+          <function>actionexecute</function>
         </task>
       </execute-stage>
+
+      <execute-stage title="Save Cluster State" service="" component="">
+        <task xsi:type="server_action" class="org.apache.ambari.server.serveraction.upgrades.FinalizeUpgradeAction" />
+      </execute-stage>
+
     </group>
   </order>
 
@@ -323,6 +331,7 @@
 
     <service name="RANGER">
       <component name="RANGER_ADMIN">
+        <pre-downgrade /> <!--  no-op to prevent config changes on downgrade -->
         <pre-upgrade>
           <task xsi:type="configure">
             <type>ranger-env</type>
@@ -432,6 +441,7 @@
 
     <service name="HDFS">
       <component name="NAMENODE">
+        <pre-downgrade /> <!--  no-op to prevent config changes on downgrade -->
         <pre-upgrade>
           <task xsi:type="configure">
             <condition type="ranger-hdfs-plugin-properties" key="ranger-hdfs-plugin-enabled" value="true">
@@ -529,6 +539,7 @@
 
     <service name="MAPREDUCE2">
       <component name="HISTORYSERVER">
+        <pre-downgrade /> <!--  no-op to prevent config changes on downgrade -->
         <pre-upgrade>
           <task xsi:type="configure">
             <type>mapred-site</type>
@@ -552,6 +563,7 @@
 
     <service name="YARN">
       <component name="APP_TIMELINE_SERVER">
+        <pre-downgrade /> <!--  no-op to prevent config changes on downgrade -->
         <pre-upgrade>
           <task xsi:type="configure">
             <type>yarn-site</type>
@@ -567,6 +579,7 @@
       </component>
 
       <component name="RESOURCEMANAGER">
+        <pre-downgrade /> <!--  no-op to prevent config changes on downgrade -->
         <pre-upgrade>
           <task xsi:type="configure">
             <type>yarn-site</type>
@@ -615,6 +628,7 @@
 
     <service name="HBASE">
       <component name="HBASE_MASTER">
+        <pre-downgrade /> <!--  no-op to prevent config changes on downgrade -->
         <pre-upgrade>
           <task xsi:type="configure" summary="Transitioning Ranger HBase Policy">
             <type>ranger-hbase-policymgr-ssl</type>
@@ -890,6 +904,7 @@
 
     <service name="OOZIE">
       <component name="OOZIE_SERVER">
+        <pre-downgrade /> <!--  no-op to prevent config changes on downgrade -->
         <pre-upgrade>
           <task xsi:type="configure">
             <summary>Updating oozie-site to remove redundant configurations</summary>
@@ -957,6 +972,7 @@
 
     <service name="KNOX">
       <component name="KNOX_GATEWAY">
+        <pre-downgrade /> <!--  no-op to prevent config changes on downgrade -->
         <pre-upgrade>
           <task xsi:type="configure" summary="Configuring Ranger Knox Policy">
             <type>ranger-knox-policymgr-ssl</type>
@@ -1021,6 +1037,22 @@
 
     <service name="STORM">
       <component name="NIMBUS">
+        <pre-downgrade>
+          <task xsi:type="manual">
+            <message>Before continuing, please deactivate and kill any currently running topologies.</message>
+          </task>
+        
+          <task xsi:type="execute" summary="Removing Storm data from ZooKeeper">
+            <script>scripts/storm_upgrade.py</script>
+            <function>delete_storm_zookeeper_data</function>
+          </task>
+
+          <task xsi:type="execute" summary="Removing local Storm data">
+            <script>scripts/storm_upgrade.py</script>
+            <function>delete_storm_local_data</function>
+          </task>
+        </pre-downgrade>
+      
         <pre-upgrade>
           <task xsi:type="manual">
             <message>Before continuing, please deactivate and kill any currently running topologies.</message>

ambari-server/src/main/resources/stacks/HDP/2.3/upgrades/upgrade-2.3.xml

@@ -321,6 +321,14 @@
           <function>finalize_rolling_upgrade</function>
         </task>
       </execute-stage>
+      
+      <execute-stage title="Update remaining HDP stack to {{version}}">
+        <task xsi:type="execute">
+          <script>scripts/ru_set_all.py</script>
+          <function>actionexecute</function>
+        </task>
+      </execute-stage>
+      
       <execute-stage title="Save Cluster State" service="" component="">
         <task xsi:type="server_action" class="org.apache.ambari.server.serveraction.upgrades.FinalizeUpgradeAction">
         </task>

ambari-server/src/test/java/org/apache/ambari/server/state/UpgradeHelperTest.java

@@ -595,6 +595,84 @@ public class UpgradeHelperTest {
         manualTask.message);
   }
 
+  @Test
+  public void testUpgradeOrchestrationFullTask() throws Exception {
+    Map<String, UpgradePack> upgrades = ambariMetaInfo.getUpgradePacks("HDP", "2.1.1");
+
+    ServiceInfo si = ambariMetaInfo.getService("HDP", "2.1.1", "ZOOKEEPER");
+    si.setDisplayName("Zk");
+    ComponentInfo ci = si.getComponentByName("ZOOKEEPER_SERVER");
+    ci.setDisplayName("ZooKeeper1 Server2");
+
+    assertTrue(upgrades.containsKey("upgrade_to_new_stack"));
+    UpgradePack upgrade = upgrades.get("upgrade_to_new_stack");
+    assertNotNull(upgrade);
+
+    makeCluster();
+
+    UpgradeContext context = new UpgradeContext(m_masterHostResolver, HDP_21,
+        HDP_21, UPGRADE_VERSION, Direction.UPGRADE);
+
+    List<UpgradeGroupHolder> groups = m_upgradeHelper.createSequence(upgrade, context);
+
+    assertEquals(6, groups.size());
+
+    assertEquals("PRE_CLUSTER", groups.get(0).name);
+    assertEquals("ZOOKEEPER", groups.get(1).name);
+    assertEquals("CORE_MASTER", groups.get(2).name);
+    assertEquals("CORE_SLAVES", groups.get(3).name);
+    assertEquals("HIVE", groups.get(4).name);
+
+    UpgradeGroupHolder holder = groups.get(2);
+    boolean found = false;
+    for (StageWrapper sw : holder.items) {
+      if (sw.getTasksJson().contains("Upgrading your database")) {
+        found = true;
+      }
+    }
+    assertTrue("Expected to find replaced text for Upgrading", found);
+
+    UpgradeGroupHolder group = groups.get(1);
+    // check that the display name is being used
+    assertTrue(group.items.get(1).getText().contains("ZooKeeper1 Server2"));
+    assertEquals(group.items.get(5).getText(), "Service Check Zk");
+
+    group = groups.get(3);
+    assertEquals(8, group.items.size());
+    StageWrapper sw = group.items.get(3);
+    assertEquals("Validate Partial Upgrade", sw.getText());
+    assertEquals(1, sw.getTasks().size());
+    assertEquals(1, sw.getTasks().get(0).getTasks().size());
+    Task t = sw.getTasks().get(0).getTasks().get(0);
+    assertEquals(ManualTask.class, t.getClass());
+    ManualTask mt = (ManualTask) t;
+    assertTrue(mt.message.contains("DataNode and NodeManager"));
+    assertNotNull(mt.structuredOut);
+    assertTrue(mt.structuredOut.contains("DATANODE"));
+    assertTrue(mt.structuredOut.contains("NODEMANAGER"));
+
+    UpgradeGroupHolder postGroup = groups.get(5);
+    assertEquals(postGroup.name, "POST_CLUSTER");
+    assertEquals(postGroup.title, "Finalize Upgrade");
+    assertEquals(4, postGroup.items.size());
+    assertEquals("Confirm Finalize", postGroup.items.get(0).getText());
+    assertEquals("Execute HDFS Finalize", postGroup.items.get(1).getText());
+    assertEquals("Save Cluster State", postGroup.items.get(2).getText());
+    assertEquals(StageWrapper.Type.SERVER_SIDE_ACTION, postGroup.items.get(2).getType());
+    assertEquals("Run On All 2.2.1.0-1234", postGroup.items.get(3).getText());
+
+    assertEquals(1, postGroup.items.get(3).getTasks().size());
+    Set<String> hosts = postGroup.items.get(3).getTasks().get(0).getHosts();
+    assertNotNull(hosts);
+    assertEquals(4, hosts.size());
+
+    assertEquals(4, groups.get(0).items.size());
+    assertEquals(6, groups.get(1).items.size());
+    assertEquals(8, groups.get(2).items.size());
+    assertEquals(8, groups.get(3).items.size());
+  }
+
+
   private Cluster makeCluster() throws AmbariException {
     return makeCluster(true);
   }
@@ -727,6 +805,8 @@ public class UpgradeHelperTest {
     return c;
   }
 
+
+
   /**
    *
    */

ambari-server/src/test/python/custom_actions/test_ru_execute_tasks.py

@@ -31,6 +31,7 @@ from resource_management import Script, ConfigDictionary
 from resource_management.libraries.functions.default import default
 from resource_management.core.logger import Logger
 from ambari_agent.AmbariConfig import AmbariConfig
+from ambari_agent.FileCache import FileCache
 from ambari_commons.os_check import OSCheck
 
 
@@ -68,7 +69,8 @@ class TestRUExecuteTasks(RMFTestCase):
   @patch.object(AmbariConfig, "getConfigFile")
   @patch.object(Script, 'get_tmp_dir')
   @patch.object(Script, 'get_config')
-  def test_execution(self, get_config_mock, get_tmp_dir_mock, get_config_file_mock, os_path_exists_mock, call_mock):
+  @patch.object(FileCache, 'get_service_base_dir')
+  def test_execution(self, cache_mock, get_config_mock, get_tmp_dir_mock, get_config_file_mock, os_path_exists_mock, call_mock):
     # Mock the config objects
     json_file_path = os.path.join(self.CUSTOM_ACTIONS_DIR, "ru_execute_tasks_namenode_prepare.json")
     self.assertTrue(os.path.isfile(json_file_path))
@@ -77,6 +79,7 @@ class TestRUExecuteTasks(RMFTestCase):
 
     config_dict = ConfigDictionary(json_payload)
 
+    cache_mock.return_value = "/var/lib/ambari-agent/cache/common-services/HDFS/2.1.0.2.0/package"
     get_config_mock.return_value = config_dict
     get_tmp_dir_mock.return_value = "/tmp"
 
@@ -101,4 +104,49 @@ class TestRUExecuteTasks(RMFTestCase):
     ru_execute = ExecuteUpgradeTasks()
     ru_execute.actionexecute(None)
 
-    call_mock.assert_called_with("/usr/bin/ambari-python-wrap /var/lib/ambari-agent/cache/common-services/HDFS/2.1.0.2.0/package/scripts/namenode.py prepare_rolling_upgrade /tmp")
+    call_mock.assert_called_with("/usr/bin/ambari-python-wrap /var/lib/ambari-agent/cache/common-services/HDFS/2.1.0.2.0/package/scripts/namenode.py prepare_rolling_upgrade /tmp")
+
+  @patch("resource_management.core.shell.call")
+  @patch("os.path.exists")
+  @patch.object(AmbariConfig, "getConfigFile")
+  @patch.object(Script, 'get_tmp_dir')
+  @patch.object(Script, 'get_config')
+  @patch.object(FileCache, 'get_custom_actions_base_dir')
+  def test_execution_custom_action(self, cache_mock, get_config_mock, get_tmp_dir_mock, get_config_file_mock, os_path_exists_mock, call_mock):
+    # Mock the config objects
+    json_file_path = os.path.join(self.CUSTOM_ACTIONS_DIR, "ru_execute_tasks_namenode_prepare.json")
+    self.assertTrue(os.path.isfile(json_file_path))
+    with open(json_file_path, "r") as json_file:
+      json_payload = json.load(json_file)
+
+    del json_payload['roleParams']['service_package_folder']
+    del json_payload['roleParams']['hooks_folder']
+
+    config_dict = ConfigDictionary(json_payload)
+
+    cache_mock.return_value = "/var/lib/ambari-agent/cache/custom_actions"
+    get_config_mock.return_value = config_dict
+    get_tmp_dir_mock.return_value = "/tmp"
+
+    ambari_agent_ini_file_path = os.path.join(os.getcwd(), "../../../../ambari-agent/conf/unix/ambari-agent.ini")
+    self.assertTrue(os.path.isfile(ambari_agent_ini_file_path))
+    get_config_file_mock.return_value = ambari_agent_ini_file_path
+
+    # Mock os calls
+    os_path_exists_mock.return_value = True
+    call_mock.side_effect = fake_call   # echo the command
+
+    # Ensure that the json file was actually read.
+    stack_name = default("/hostLevelParams/stack_name", None)
+    stack_version = default("/hostLevelParams/stack_version", None)
+    service_package_folder = default('/roleParams/service_package_folder', None)
+
+    self.assertEqual(stack_name, "HDP")
+    self.assertEqual(stack_version, 2.2)
+    self.assertEqual(service_package_folder, None)
+
+    # Begin the test
+    ru_execute = ExecuteUpgradeTasks()
+    ru_execute.actionexecute(None)
+
+    call_mock.assert_called_with("/usr/bin/ambari-python-wrap /var/lib/ambari-agent/cache/custom_actions/scripts/namenode.py prepare_rolling_upgrade /tmp")

ambari-server/src/test/python/custom_actions/test_ru_set_all.py

@@ -0,0 +1,94 @@
+# !/usr/bin/env python
+
+'''
+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.
+'''
+
+# Python Imports
+import os
+import json
+
+from mock.mock import patch
+from mock.mock import MagicMock
+
+# Module imports
+from stacks.utils.RMFTestCase import *
+from resource_management import Script, ConfigDictionary
+from resource_management.libraries.functions.default import default
+from resource_management.core.logger import Logger
+from ambari_agent.AmbariConfig import AmbariConfig
+from ambari_agent.FileCache import FileCache
+from ambari_commons.os_check import OSCheck
+
+
+def fake_call(command):
+  """
+  Instead of shell.call, call a command whose output equals the command.
+  :param command: Command that will be echoed.
+  :return: Returns a tuple of (process output code, output)
+  """
+  return (0, command)
+
+
+class TestRUSetAll(RMFTestCase):
+  CUSTOM_ACTIONS_DIR = os.path.join(os.getcwd(), "../resources/custom_actions/")
+
+  @patch.object(Logger, "info")
+  @patch.object(Logger, "error")
+  def setUp(self, error_mock, info_mock):
+
+    Logger.logger = MagicMock()
+
+    # Import the class under test. This is done here as opposed to the rest of the imports because the get_os_type()
+    # method needs to be patched first.
+    from ru_set_all import UpgradeSetAll
+    global UpgradeSetAll
+
+  def tearDown(self):
+    Logger.logger = None
+
+  @patch("resource_management.core.shell.call")
+  @patch.object(Script, 'get_config')
+  @patch.object(OSCheck, 'is_redhat_family')
+  def test_execution(self, family_mock, get_config_mock, call_mock):
+    # Mock the config objects
+    json_file_path = os.path.join(self.CUSTOM_ACTIONS_DIR, "ru_execute_tasks_namenode_prepare.json")
+    self.assertTrue(os.path.isfile(json_file_path))
+    with open(json_file_path, "r") as json_file:
+      json_payload = json.load(json_file)
+
+    config_dict = ConfigDictionary(json_payload)
+
+    family_mock.return_value = True
+    get_config_mock.return_value = config_dict
+    call_mock.side_effect = fake_call   # echo the command
+
+    # Ensure that the json file was actually read.
+    stack_name = default("/hostLevelParams/stack_name", None)
+    stack_version = default("/hostLevelParams/stack_version", None)
+    service_package_folder = default('/roleParams/service_package_folder', None)
+
+    self.assertEqual(stack_name, "HDP")
+    self.assertEqual(stack_version, 2.2)
+    self.assertEqual(service_package_folder, "common-services/HDFS/2.1.0.2.0/package")
+
+    # Begin the test
+    ru_execute = UpgradeSetAll()
+    ru_execute.actionexecute(None)
+
+    call_mock.assert_called_with("hdp-select set all 2.2.1.0-2260")
+

ambari-server/src/test/resources/stacks/HDP/2.1.1/upgrades/upgrade_to_new_stack.xml

@@ -97,20 +97,29 @@
               {{host-detail-list}}</message>
         </task>
       </execute-stage>
+      
       <execute-stage title="Confirm Finalize">
         <task xsi:type="manual">
           <message>Please confirm you are ready to finalize</message>
         </task>
       </execute-stage>
+      
       <execute-stage service="HDFS" component="NAMENODE" title="Execute HDFS Finalize">
         <task xsi:type="execute">
           <command>ls</command>
         </task>
       </execute-stage>
+      
       <execute-stage title="Save Cluster State" service="" component="">
         <task xsi:type="server_action" class="org.apache.ambari.server.serveraction.upgrades.FinalizeUpgradeAction">
         </task>
       </execute-stage>
+      
+      <execute-stage title="Run On All {{version}}">
+        <task xsi:type="execute">
+          <command>ls</command>
+        </task>
+      </execute-stage>
     </group>
         
   </order>