@Test
  public void testOrchestrationOptions() throws Exception {

    Map<String, UpgradePack> upgrades = ambariMetaInfo.getUpgradePacks("HDP", "2.2.0");
    assertTrue(upgrades.containsKey("upgrade_from_211"));

    UpgradePack upgrade = upgrades.get("upgrade_from_211");
    assertNotNull(upgrade);

    Cluster cluster = makeCluster();

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

    SimpleOrchestrationOptions options = new SimpleOrchestrationOptions(1);

    expect(context.getOrchestrationOptions()).andReturn(options).anyTimes();
    replay(context);

    List<UpgradeGroupHolder> groups = m_upgradeHelper.createSequence(upgrade, context);
    groups = groups.stream().filter(g -> g.name.equals("CORE_SLAVES")).collect(Collectors.toList());
    assertEquals(1, groups.size());

    List<StageWrapper> restarts = groups.get(0).items.stream().filter(sw ->
        sw.getType() == StageWrapper.Type.RESTART && sw.getText().contains("DataNode"))
        .collect(Collectors.toList());

    assertEquals("Expecting wrappers for each of 3 hosts", 3, restarts.size());

    options.m_count = 2;
    groups = m_upgradeHelper.createSequence(upgrade, context);
    groups = groups.stream().filter(g -> g.name.equals("CORE_SLAVES")).collect(Collectors.toList());
    assertEquals(1, groups.size());

    restarts = groups.get(0).items.stream().filter(sw ->
        sw.getType() == StageWrapper.Type.RESTART && sw.getText().contains("DataNode"))
        .collect(Collectors.toList());

    assertEquals("Expecting wrappers for each", 2, restarts.size());
  }
Tests that components added during the upgrade are scheduled for restart on
their future hosts.

@throws Exception
ambari-server/src/test/java/org/apache/ambari/server/stack/upgrade/orchestrate/UpgradeHelperTest.java