@Test
  public void testUpgradeWithMultipleTasksInOwnStage() throws Exception {
    Map<String, UpgradePack> upgrades = ambariMetaInfo.getUpgradePacks("HDP", "2.1.1");

    assertTrue(upgrades.containsKey("upgrade_test"));
    UpgradePack upgrade = upgrades.get("upgrade_test");
    assertNotNull(upgrade);
    assertTrue(upgrade.getType() == UpgradeType.ROLLING);

    List<Grouping> upgradePackGroups = upgrade.getGroups(Direction.UPGRADE);

    boolean foundService = false;
    for (Grouping group : upgradePackGroups) {
      if (group.title.equals("Oozie")) {
        foundService = true;
      }
    }
    assertTrue(foundService);

    Cluster cluster = makeCluster();

    UpgradeContext context = getMockUpgradeContext(cluster, Direction.UPGRADE, UpgradeType.ROLLING);

    List<UpgradeGroupHolder> groups = m_upgradeHelper.createSequence(upgrade, context);

    // The upgrade pack has 2 tasks for Oozie in the pre-upgrade group.
    // The first task runs on "all", i.e., both Oozie Servers, whereas the
    // second task runs on "any", i.e., exactly one.
    int numPrepareStages = 0;
    for (UpgradeGroupHolder group : groups) {
      if (group.name.equals("OOZIE")) {
        assertTrue(group.items.size() > 0);
        for (StageWrapper sw : group.items) {

          if (sw.getText().equalsIgnoreCase("Preparing Oozie Server on h2 (Batch 1 of 2)") ||
              sw.getText().equalsIgnoreCase("Preparing Oozie Server on h3 (Batch 2 of 2)")) {
            numPrepareStages++;
            List<TaskWrapper> taskWrappers = sw.getTasks();
            assertEquals(1, taskWrappers.size());
            List<Task> tasks = taskWrappers.get(0).getTasks();
            assertEquals(1, taskWrappers.get(0).getHosts().size());
            assertEquals(1, tasks.size());

            ExecuteTask task = (ExecuteTask) tasks.get(0);
            assertTrue("scripts/oozie_server.py".equalsIgnoreCase(task.script));
            assertTrue("stop".equalsIgnoreCase(task.function));
          }

          if (sw.getText().equalsIgnoreCase("Preparing Oozie Server on h2")) {
            numPrepareStages++;
            List<TaskWrapper> taskWrappers = sw.getTasks();
            assertEquals(1, taskWrappers.size());
            List<Task> tasks = taskWrappers.get(0).getTasks();
            assertEquals(1, taskWrappers.get(0).getHosts().size());
            assertEquals(1, tasks.size());

            ExecuteTask task = (ExecuteTask) tasks.get(0);
            assertTrue("scripts/oozie_server_upgrade.py".equalsIgnoreCase(task.script));
            assertTrue("upgrade_oozie_database_and_sharelib".equalsIgnoreCase(task.function));
          }
        }
      }
    }
    assertEquals(3, numPrepareStages);
  }
Test that multiple execute tasks with an annotation of synchronized="true" each run in their own stage.
ambari-server/src/test/java/org/apache/ambari/server/stack/upgrade/orchestrate/UpgradeHelperTest.java