You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@ambari.apache.org by jo...@apache.org on 2017/06/28 18:07:44 UTC

ambari git commit: AMBARI-21367 - Fix Unit Test Failures Caused by AMBARI-21358 (jonathanhurley)

Repository: ambari
Updated Branches:
  refs/heads/branch-feature-AMBARI-21348 1863c3b90 -> 77099797a


AMBARI-21367 - Fix Unit Test Failures Caused by AMBARI-21358 (jonathanhurley)


Project: http://git-wip-us.apache.org/repos/asf/ambari/repo
Commit: http://git-wip-us.apache.org/repos/asf/ambari/commit/77099797
Tree: http://git-wip-us.apache.org/repos/asf/ambari/tree/77099797
Diff: http://git-wip-us.apache.org/repos/asf/ambari/diff/77099797

Branch: refs/heads/branch-feature-AMBARI-21348
Commit: 77099797a6326d285245e8af8dd156900804cf31
Parents: 1863c3b
Author: Jonathan Hurley <jh...@hortonworks.com>
Authored: Wed Jun 28 12:46:06 2017 -0400
Committer: Jonathan Hurley <jh...@hortonworks.com>
Committed: Wed Jun 28 13:53:39 2017 -0400

----------------------------------------------------------------------
 .../actionmanager/ExecutionCommandWrapper.java  |   6 +-
 .../internal/UpgradeResourceProvider.java       |  13 +-
 .../upgrades/UpdateDesiredStackAction.java      |   6 +-
 .../internal/UpgradeResourceProviderTest.java   |  18 +--
 .../UpgradeSummaryResourceProviderTest.java     |   8 +-
 .../ComponentVersionCheckActionTest.java        | 120 +++++++++++-----
 .../upgrades/UpgradeActionTest.java             | 136 ++++++++++---------
 .../ambari/server/state/UpgradeHelperTest.java  | 127 +++++++++++++----
 .../cluster/ClusterEffectiveVersionTest.java    |   2 +-
 9 files changed, 295 insertions(+), 141 deletions(-)
----------------------------------------------------------------------


http://git-wip-us.apache.org/repos/asf/ambari/blob/77099797/ambari-server/src/main/java/org/apache/ambari/server/actionmanager/ExecutionCommandWrapper.java
----------------------------------------------------------------------
diff --git a/ambari-server/src/main/java/org/apache/ambari/server/actionmanager/ExecutionCommandWrapper.java b/ambari-server/src/main/java/org/apache/ambari/server/actionmanager/ExecutionCommandWrapper.java
index 81bc80f..4773c75 100644
--- a/ambari-server/src/main/java/org/apache/ambari/server/actionmanager/ExecutionCommandWrapper.java
+++ b/ambari-server/src/main/java/org/apache/ambari/server/actionmanager/ExecutionCommandWrapper.java
@@ -183,8 +183,10 @@ public class ExecutionCommandWrapper {
       }
 
       ClusterVersionEntity effectiveClusterVersion = cluster.getEffectiveClusterVersion();
-      executionCommand.getCommandParams().put(KeyNames.VERSION,
-          effectiveClusterVersion.getRepositoryVersion().getVersion());
+      if (null != effectiveClusterVersion) {
+        executionCommand.getCommandParams().put(KeyNames.VERSION,
+            effectiveClusterVersion.getRepositoryVersion().getVersion());
+      }
 
     } catch (ClusterNotFoundException cnfe) {
       // it's possible that there are commands without clusters; in such cases,

http://git-wip-us.apache.org/repos/asf/ambari/blob/77099797/ambari-server/src/main/java/org/apache/ambari/server/controller/internal/UpgradeResourceProvider.java
----------------------------------------------------------------------
diff --git a/ambari-server/src/main/java/org/apache/ambari/server/controller/internal/UpgradeResourceProvider.java b/ambari-server/src/main/java/org/apache/ambari/server/controller/internal/UpgradeResourceProvider.java
index da716b3..d2573e1 100644
--- a/ambari-server/src/main/java/org/apache/ambari/server/controller/internal/UpgradeResourceProvider.java
+++ b/ambari-server/src/main/java/org/apache/ambari/server/controller/internal/UpgradeResourceProvider.java
@@ -375,14 +375,14 @@ public class UpgradeResourceProvider extends AbstractControllerResourceProvider
 
         // use a combination of the supplied information to find the target
         // repository
-        final Long repositoryId = (Long) requestMap.get(UPGRADE_REPO_ID);
+        final String repositoryId = (String) requestMap.get(UPGRADE_REPO_ID);
         final String repositoryName = (String) requestMap.get(UPGRADE_REPO_NAME);
         final String repositoryVersion = (String) requestMap.get(UPGRADE_REPO_VERSION);
 
         RepositoryVersionEntity toRepositoryVersion = null;
 
         if (null != repositoryId) {
-          toRepositoryVersion = s_repoVersionDAO.findByPK(repositoryId);
+          toRepositoryVersion = s_repoVersionDAO.findByPK(Long.valueOf(repositoryId));
         } else {
           if (StringUtils.isNotBlank(repositoryName)) {
             toRepositoryVersion = s_repoVersionDAO.findByStackNameAndVersion(repositoryName,
@@ -665,6 +665,7 @@ public class UpgradeResourceProvider extends AbstractControllerResourceProvider
     Direction direction = upgradeContext.getDirection();
     Map<String, Object> requestMap = upgradeContext.getUpgradeRequest();
     UpgradeType upgradeType = upgradeContext.getType();
+    RepositoryVersionEntity targetRepositoryVersion = upgradeContext.getTargetRepositoryVersion();
 
     /**
      * For the unit tests tests, there are multiple upgrade packs for the same type, so
@@ -672,7 +673,7 @@ public class UpgradeResourceProvider extends AbstractControllerResourceProvider
      */
     String preferredUpgradePackName = (String) requestMap.get(UPGRADE_PACK);
 
-    String version = (String) requestMap.get(UPGRADE_REPO_VERSION);
+    String version = targetRepositoryVersion.getVersion();
     String versionForUpgradePack = (String) requestMap.get(UPGRADE_FROM_VERSION);
 
     UpgradePack pack = s_upgradeHelper.suggestUpgradePack(cluster.getClusterName(),
@@ -1967,6 +1968,7 @@ public class UpgradeResourceProvider extends AbstractControllerResourceProvider
       Map<String, Object> requestMap = upgradeContext.getUpgradeRequest();
 
       String clusterName = (String) requestMap.get(UPGRADE_CLUSTER_NAME);
+      String versionId = (String) requestMap.get(UPGRADE_REPO_ID);
       String version = (String) requestMap.get(UPGRADE_REPO_VERSION);
       String direction = (String) requestMap.get(UPGRADE_DIRECTION);
 
@@ -1974,8 +1976,9 @@ public class UpgradeResourceProvider extends AbstractControllerResourceProvider
         throw new AmbariException(String.format("%s is required", UPGRADE_CLUSTER_NAME));
       }
 
-      if (StringUtils.isBlank(version)) {
-        throw new AmbariException(String.format("%s is required", UPGRADE_REPO_VERSION));
+      if (StringUtils.isBlank(versionId) && StringUtils.isBlank(version)) {
+        throw new AmbariException(
+            String.format("Either %s or %s is required", UPGRADE_REPO_ID, UPGRADE_REPO_VERSION));
       }
 
       if (StringUtils.isBlank(direction)) {

http://git-wip-us.apache.org/repos/asf/ambari/blob/77099797/ambari-server/src/main/java/org/apache/ambari/server/serveraction/upgrades/UpdateDesiredStackAction.java
----------------------------------------------------------------------
diff --git a/ambari-server/src/main/java/org/apache/ambari/server/serveraction/upgrades/UpdateDesiredStackAction.java b/ambari-server/src/main/java/org/apache/ambari/server/serveraction/upgrades/UpdateDesiredStackAction.java
index e3f9824..cca87bc 100644
--- a/ambari-server/src/main/java/org/apache/ambari/server/serveraction/upgrades/UpdateDesiredStackAction.java
+++ b/ambari-server/src/main/java/org/apache/ambari/server/serveraction/upgrades/UpdateDesiredStackAction.java
@@ -36,7 +36,6 @@ import org.apache.ambari.server.state.StackId;
 import org.apache.ambari.server.state.StackInfo;
 import org.apache.ambari.server.state.UpgradeContext;
 import org.apache.ambari.server.state.UpgradeContextFactory;
-import org.apache.ambari.server.state.stack.UpgradePack;
 import org.apache.ambari.server.state.stack.upgrade.Direction;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
@@ -99,7 +98,6 @@ public class UpdateDesiredStackAction extends AbstractServerAction {
   private CommandReport updateDesiredStack(Cluster cluster, UpgradeContext context)
       throws AmbariException, InterruptedException {
 
-    UpgradePack upgradePack = context.getUpgradePack();
     Direction direction = context.getDirection();
     RepositoryVersionEntity fromRepositoryVersion = context.getSourceRepositoryVersion();
     RepositoryVersionEntity toRepositoryVersion = context.getTargetRepositoryVersion();
@@ -110,12 +108,12 @@ public class UpdateDesiredStackAction extends AbstractServerAction {
 
     try {
       StackId currentClusterStackId = cluster.getCurrentStackVersion();
-      out.append(String.format("%s %s from %s-%s to %s-%s using %s\n",
+      out.append(String.format("%s %s from %s-%s to %s-%s\n",
           direction.getVerb(true), clusterName,
           fromRepositoryVersion.getStackId().getStackName(),
           fromRepositoryVersion.getVersion(),
           toRepositoryVersion.getStackId().getStackName(),
-          toRepositoryVersion.getVersion(), upgradePack.getName()));
+          toRepositoryVersion.getVersion()));
 
       out.append(String.format(
           "Checking if can update the desired stack to %s. The cluster's current stack is %s\n",

http://git-wip-us.apache.org/repos/asf/ambari/blob/77099797/ambari-server/src/test/java/org/apache/ambari/server/controller/internal/UpgradeResourceProviderTest.java
----------------------------------------------------------------------
diff --git a/ambari-server/src/test/java/org/apache/ambari/server/controller/internal/UpgradeResourceProviderTest.java b/ambari-server/src/test/java/org/apache/ambari/server/controller/internal/UpgradeResourceProviderTest.java
index 9652c69..3a26164 100644
--- a/ambari-server/src/test/java/org/apache/ambari/server/controller/internal/UpgradeResourceProviderTest.java
+++ b/ambari-server/src/test/java/org/apache/ambari/server/controller/internal/UpgradeResourceProviderTest.java
@@ -642,7 +642,7 @@ public class UpgradeResourceProviderTest {
 
     Map<String, Object> requestProps = new HashMap<>();
     requestProps.put(UpgradeResourceProvider.UPGRADE_CLUSTER_NAME, "c1");
-    requestProps.put(UpgradeResourceProvider.UPGRADE_REPO_ID, repoVersionEntity2111.getId());
+    requestProps.put(UpgradeResourceProvider.UPGRADE_REPO_ID, String.valueOf(repoVersionEntity2111.getId()));
     requestProps.put(UpgradeResourceProvider.UPGRADE_PACK, "upgrade_test");
     requestProps.put(UpgradeResourceProvider.UPGRADE_SKIP_PREREQUISITE_CHECKS, "true");
     requestProps.put(UpgradeResourceProvider.UPGRADE_DIRECTION, Direction.DOWNGRADE.name());
@@ -710,10 +710,10 @@ public class UpgradeResourceProviderTest {
 
     // fix the properties and try again
     requestProps.put(UpgradeResourceProvider.UPGRADE_CLUSTER_NAME, "c1");
-    requestProps.put(UpgradeResourceProvider.UPGRADE_REPO_VERSION, "2.2.0.0");
+    requestProps.put(UpgradeResourceProvider.UPGRADE_REPO_VERSION, "2.1.1.0");
     requestProps.put(UpgradeResourceProvider.UPGRADE_PACK, "upgrade_test");
     requestProps.put(UpgradeResourceProvider.UPGRADE_SKIP_PREREQUISITE_CHECKS, "true");
-    requestProps.put(UpgradeResourceProvider.UPGRADE_FROM_VERSION, "2.1.1.0");
+    requestProps.put(UpgradeResourceProvider.UPGRADE_FROM_VERSION, "2.1.1.1");
     requestProps.put(UpgradeResourceProvider.UPGRADE_DIRECTION, Direction.DOWNGRADE.name());
 
     Map<String, String> requestInfoProperties = new HashMap<>();
@@ -726,8 +726,8 @@ public class UpgradeResourceProviderTest {
 
     UpgradeEntity entity = upgradeDao.findUpgrade(Long.parseLong(id));
     assertNotNull(entity);
-    assertEquals("2.1.1.0", entity.getFromRepositoryVersion().getVersion());
-    assertEquals("2.2.0.0", entity.getToRepositoryVersion().getVersion());
+    assertEquals("2.1.1.1", entity.getFromRepositoryVersion().getVersion());
+    assertEquals("2.1.1.0", entity.getToRepositoryVersion().getVersion());
     assertEquals(Direction.DOWNGRADE, entity.getDirection());
 
     StageDAO dao = injector.getInstance(StageDAO.class);
@@ -1201,10 +1201,6 @@ public class UpgradeResourceProviderTest {
 
     EasyMock.expectLastCall().once();
 
-    EasyMock.replay(configHelper, cluster, fooConfig, barConfig, bazConfig);
-
-    UpgradeResourceProvider upgradeResourceProvider = createProvider(amc);
-
     Map<String, UpgradePack> upgradePacks = ambariMetaInfo.getUpgradePacks("HDP", "2.1.1");
     UpgradePack upgradePack = upgradePacks.get("upgrade_to_new_stack");
 
@@ -1216,6 +1212,10 @@ public class UpgradeResourceProviderTest {
     expect(context.getTargetRepositoryVersion()).andReturn(repoVersionEntity2200).atLeastOnce();
     expect(context.getUpgradePack()).andReturn(upgradePack).atLeastOnce();
 
+    EasyMock.replay(configHelper, cluster, fooConfig, barConfig, bazConfig, context);
+
+    UpgradeResourceProvider upgradeResourceProvider = createProvider(amc);
+
     upgradeResourceProvider.applyStackAndProcessConfigurations(context);
 
     Map<String, Map<String, String>> expectedConfigurations = expectedConfigurationsCapture.getValue();

http://git-wip-us.apache.org/repos/asf/ambari/blob/77099797/ambari-server/src/test/java/org/apache/ambari/server/controller/internal/UpgradeSummaryResourceProviderTest.java
----------------------------------------------------------------------
diff --git a/ambari-server/src/test/java/org/apache/ambari/server/controller/internal/UpgradeSummaryResourceProviderTest.java b/ambari-server/src/test/java/org/apache/ambari/server/controller/internal/UpgradeSummaryResourceProviderTest.java
index f63825b..2c6f61f 100644
--- a/ambari-server/src/test/java/org/apache/ambari/server/controller/internal/UpgradeSummaryResourceProviderTest.java
+++ b/ambari-server/src/test/java/org/apache/ambari/server/controller/internal/UpgradeSummaryResourceProviderTest.java
@@ -155,7 +155,7 @@ public class UpgradeSummaryResourceProviderTest {
     repoVersionEntity.setDisplayName("For Stack Version 2.2.0");
     repoVersionEntity.setOperatingSystems("");
     repoVersionEntity.setStack(stackEntity);
-    repoVersionEntity.setVersion("2.2.0.0");
+    repoVersionEntity.setVersion("2.2.0.0-1234");
     repoVersionDAO.create(repoVersionEntity);
 
     clusters = injector.getInstance(Clusters.class);
@@ -164,7 +164,7 @@ public class UpgradeSummaryResourceProviderTest {
     clusters.addCluster(clusterName, stackId);
     Cluster cluster = clusters.getCluster("c1");
 
-    helper.getOrCreateRepositoryVersion(stackId, stackId.getStackVersion());
+    helper.getOrCreateRepositoryVersion(stackId, "2.2.0.1-1234");
     cluster.createClusterVersion(stackId, stackId.getStackVersion(), "admin", RepositoryVersionState.INSTALLING);
 
     clusters.addHost("h1");
@@ -183,11 +183,11 @@ public class UpgradeSummaryResourceProviderTest {
 
     ServiceComponent component = service.addServiceComponent("ZOOKEEPER_SERVER");
     ServiceComponentHost sch = component.addServiceComponentHost("h1");
-    sch.setVersion("2.2.0.0");
+    sch.setVersion("2.2.0.0-1234");
 
     component = service.addServiceComponent("ZOOKEEPER_CLIENT");
     sch = component.addServiceComponentHost("h1");
-    sch.setVersion("2.2.0.0");
+    sch.setVersion("2.2.0.0-1234");
   }
 
   /**

http://git-wip-us.apache.org/repos/asf/ambari/blob/77099797/ambari-server/src/test/java/org/apache/ambari/server/serveraction/upgrades/ComponentVersionCheckActionTest.java
----------------------------------------------------------------------
diff --git a/ambari-server/src/test/java/org/apache/ambari/server/serveraction/upgrades/ComponentVersionCheckActionTest.java b/ambari-server/src/test/java/org/apache/ambari/server/serveraction/upgrades/ComponentVersionCheckActionTest.java
index ceaa83a..47d2a81 100644
--- a/ambari-server/src/test/java/org/apache/ambari/server/serveraction/upgrades/ComponentVersionCheckActionTest.java
+++ b/ambari-server/src/test/java/org/apache/ambari/server/serveraction/upgrades/ComponentVersionCheckActionTest.java
@@ -46,10 +46,15 @@ import org.apache.ambari.server.orm.dao.ClusterVersionDAO;
 import org.apache.ambari.server.orm.dao.HostDAO;
 import org.apache.ambari.server.orm.dao.HostVersionDAO;
 import org.apache.ambari.server.orm.dao.RepositoryVersionDAO;
+import org.apache.ambari.server.orm.dao.RequestDAO;
 import org.apache.ambari.server.orm.dao.StackDAO;
+import org.apache.ambari.server.orm.dao.UpgradeDAO;
 import org.apache.ambari.server.orm.entities.ClusterVersionEntity;
 import org.apache.ambari.server.orm.entities.HostVersionEntity;
+import org.apache.ambari.server.orm.entities.RepositoryVersionEntity;
+import org.apache.ambari.server.orm.entities.RequestEntity;
 import org.apache.ambari.server.orm.entities.StackEntity;
+import org.apache.ambari.server.orm.entities.UpgradeEntity;
 import org.apache.ambari.server.state.Cluster;
 import org.apache.ambari.server.state.Clusters;
 import org.apache.ambari.server.state.ConfigFactory;
@@ -64,6 +69,7 @@ import org.apache.ambari.server.state.ServiceComponentHostFactory;
 import org.apache.ambari.server.state.ServiceFactory;
 import org.apache.ambari.server.state.StackId;
 import org.apache.ambari.server.state.State;
+import org.apache.ambari.server.state.stack.upgrade.UpgradeType;
 import org.apache.ambari.server.utils.EventBusSynchronizer;
 import org.junit.After;
 import org.junit.Assert;
@@ -146,6 +152,9 @@ public class ComponentVersionCheckActionTest {
     clusters.addCluster(clusterName, sourceStack);
 
     StackDAO stackDAO = m_injector.getInstance(StackDAO.class);
+    RequestDAO requestDAO = m_injector.getInstance(RequestDAO.class);
+    UpgradeDAO upgradeDAO = m_injector.getInstance(UpgradeDAO.class);
+
     StackEntity stackEntitySource = stackDAO.find(sourceStack.getStackName(), sourceStack.getStackVersion());
     StackEntity stackEntityTarget = stackDAO.find(targetStack.getStackName(), targetStack.getStackVersion());
     assertNotNull(stackEntitySource);
@@ -159,13 +168,13 @@ public class ComponentVersionCheckActionTest {
 
     Host host = clusters.getHost(hostName);
 
-    Map<String, String> hostAttributes = new HashMap<String, String>();
+    Map<String, String> hostAttributes = new HashMap<>();
     hostAttributes.put("os_family", "redhat");
     hostAttributes.put("os_release_version", "6");
     host.setHostAttributes(hostAttributes);
 
     // Create the starting repo version
-    m_helper.getOrCreateRepositoryVersion(sourceStack, sourceRepo);
+    RepositoryVersionEntity sourceRepositoryVersion = m_helper.getOrCreateRepositoryVersion(sourceStack, sourceRepo);
     c.createClusterVersion(sourceStack, sourceRepo, "admin", RepositoryVersionState.INSTALLING);
     c.transitionClusterVersion(sourceStack, sourceRepo, RepositoryVersionState.CURRENT);
 
@@ -173,7 +182,9 @@ public class ComponentVersionCheckActionTest {
     String urlInfo = "[{'repositories':["
         + "{'Repositories/base_url':'http://foo1','Repositories/repo_name':'HDP','Repositories/repo_id':'" + targetStack.getStackId() + "'}"
         + "], 'OperatingSystems/os_type':'redhat6'}]";
-    repoVersionDAO.create(stackEntityTarget, targetRepo, String.valueOf(System.currentTimeMillis()), urlInfo);
+
+    RepositoryVersionEntity toRepositoryVersion = repoVersionDAO.create(stackEntityTarget,
+        targetRepo, String.valueOf(System.currentTimeMillis()), urlInfo);
 
     // Start upgrading the newer repo
     c.createClusterVersion(targetStack, targetRepo, "admin", RepositoryVersionState.INSTALLING);
@@ -190,6 +201,25 @@ public class ComponentVersionCheckActionTest {
     entity.setRepositoryVersion(repoVersionDAO.findByStackAndVersion(targetStack, targetRepo));
     entity.setState(RepositoryVersionState.INSTALLED);
     hostVersionDAO.create(entity);
+
+    RequestEntity requestEntity = new RequestEntity();
+    requestEntity.setClusterId(c.getClusterId());
+    requestEntity.setRequestId(1L);
+    requestEntity.setStartTime(System.currentTimeMillis());
+    requestEntity.setCreateTime(System.currentTimeMillis());
+    requestDAO.create(requestEntity);
+
+    UpgradeEntity upgradeEntity = new UpgradeEntity();
+    upgradeEntity.setId(1L);
+    upgradeEntity.setClusterId(c.getClusterId());
+    upgradeEntity.setRequestEntity(requestEntity);
+    upgradeEntity.setUpgradePackage("");
+    upgradeEntity.setFromRepositoryVersion(sourceRepositoryVersion);
+    upgradeEntity.setToRepositoryVersion(toRepositoryVersion);
+    upgradeEntity.setUpgradeType(UpgradeType.NON_ROLLING);
+    upgradeDAO.create(upgradeEntity);
+
+    c.setUpgradeEntity(upgradeEntity);
   }
 
   private void makeCrossStackUpgradeCluster(StackId sourceStack, String sourceRepo, StackId targetStack,
@@ -199,6 +229,9 @@ public class ComponentVersionCheckActionTest {
     clusters.addCluster(clusterName, sourceStack);
 
     StackDAO stackDAO = m_injector.getInstance(StackDAO.class);
+    RequestDAO requestDAO = m_injector.getInstance(RequestDAO.class);
+    UpgradeDAO upgradeDAO = m_injector.getInstance(UpgradeDAO.class);
+
     StackEntity stackEntitySource = stackDAO.find(sourceStack.getStackName(), sourceStack.getStackVersion());
     StackEntity stackEntityTarget = stackDAO.find(targetStack.getStackName(), targetStack.getStackVersion());
 
@@ -213,7 +246,7 @@ public class ComponentVersionCheckActionTest {
     clusters.addHost(hostName);
     Host host = clusters.getHost(hostName);
 
-    Map<String, String> hostAttributes = new HashMap<String, String>();
+    Map<String, String> hostAttributes = new HashMap<>();
     hostAttributes.put("os_family", "redhat");
     hostAttributes.put("os_release_version", "6");
     host.setHostAttributes(hostAttributes);
@@ -221,38 +254,49 @@ public class ComponentVersionCheckActionTest {
     clusters.mapHostToCluster(hostName, clusterName);
 
     // Create the starting repo version
-    m_helper.getOrCreateRepositoryVersion(sourceStack, sourceRepo);
+    RepositoryVersionEntity sourceRepositoryVersion = m_helper.getOrCreateRepositoryVersion(sourceStack, sourceRepo);
     c.createClusterVersion(sourceStack, sourceRepo, "admin", RepositoryVersionState.INSTALLING);
     c.transitionClusterVersion(sourceStack, sourceRepo, RepositoryVersionState.CURRENT);
 
+    RepositoryVersionEntity toRepositoryVersion = m_helper.getOrCreateRepositoryVersion(targetStack,targetRepo);
+
+    RequestEntity requestEntity = new RequestEntity();
+    requestEntity.setClusterId(c.getClusterId());
+    requestEntity.setRequestId(1L);
+    requestEntity.setStartTime(System.currentTimeMillis());
+    requestEntity.setCreateTime(System.currentTimeMillis());
+    requestDAO.create(requestEntity);
+
+    UpgradeEntity upgradeEntity = new UpgradeEntity();
+    upgradeEntity.setId(1L);
+    upgradeEntity.setClusterId(c.getClusterId());
+    upgradeEntity.setRequestEntity(requestEntity);
+    upgradeEntity.setUpgradePackage("");
+    upgradeEntity.setFromRepositoryVersion(sourceRepositoryVersion);
+    upgradeEntity.setToRepositoryVersion(toRepositoryVersion);
+    upgradeEntity.setUpgradeType(UpgradeType.NON_ROLLING);
+    upgradeDAO.create(upgradeEntity);
+
+    c.setUpgradeEntity(upgradeEntity);
   }
 
-  private void createNewRepoVersion(StackId targetStack, String targetRepo, String clusterName,
-                                    String hostName) throws AmbariException {
-    Clusters clusters = m_injector.getInstance(Clusters.class);
-    StackDAO stackDAO = m_injector.getInstance(StackDAO.class);
-
-    StackEntity stackEntityTarget = stackDAO.find(targetStack.getStackName(), targetStack.getStackVersion());
-
-    Cluster c = clusters.getCluster(clusterName);
-    // Create the new repo version
-    String urlInfo = "[{'repositories':["
-        + "{'Repositories/base_url':'http://foo1','Repositories/repo_name':'HDP','Repositories/repo_id':'" + targetRepo + "'}"
-        + "], 'OperatingSystems/os_type':'redhat6'}]";
-    repoVersionDAO.create(stackEntityTarget, targetRepo, String.valueOf(System.currentTimeMillis()), urlInfo);
 
+  /**
+   * Creates a new {@link HostVersionEntity} instance in the
+   * {@link RepositoryVersionState#INSTALLED} for the specified host.
+   *
+   * @param hostName
+   * @param repositoryVersion
+   * @throws AmbariException
+   */
+  private void installRepositoryOnHost(String hostName, RepositoryVersionEntity repositoryVersion)
+      throws AmbariException {
     // Start upgrading the newer repo
-    c.createClusterVersion(targetStack, targetRepo, "admin", RepositoryVersionState.INSTALLING);
-    c.transitionClusterVersion(targetStack, targetRepo, RepositoryVersionState.INSTALLED);
-
-    c.mapHostVersions(Collections.singleton(hostName), c.getCurrentClusterVersion(),
-        RepositoryVersionState.CURRENT);
-
     HostDAO hostDAO = m_injector.getInstance(HostDAO.class);
 
     HostVersionEntity entity = new HostVersionEntity();
     entity.setHostEntity(hostDAO.findByName(hostName));
-    entity.setRepositoryVersion(repoVersionDAO.findByStackAndVersion(targetStack, targetRepo));
+    entity.setRepositoryVersion(repositoryVersion);
     entity.setState(RepositoryVersionState.INSTALLED);
     hostVersionDAO.create(entity);
   }
@@ -273,7 +317,7 @@ public class ComponentVersionCheckActionTest {
     assertEquals(HDP_211_CENTOS6_REPO_URL, repo.getBaseUrl());
 
     // Finalize the upgrade
-    Map<String, String> commandParams = new HashMap<String, String>();
+    Map<String, String> commandParams = new HashMap<>();
     commandParams.put(FinalizeUpgradeAction.UPGRADE_DIRECTION_KEY, "upgrade");
     commandParams.put(FinalizeUpgradeAction.VERSION_KEY, targetRepo);
 
@@ -309,16 +353,30 @@ public class ComponentVersionCheckActionTest {
     Clusters clusters = m_injector.getInstance(Clusters.class);
     Cluster cluster = clusters.getCluster("c1");
 
+    RepositoryVersionEntity sourceRepoVersion = m_helper.getOrCreateRepositoryVersion(HDP_21_STACK, HDP_2_1_1_0);
+    RepositoryVersionEntity targetRepoVersion = m_helper.getOrCreateRepositoryVersion(HDP_22_STACK, HDP_2_2_1_0);
+
+    // Start upgrading the newer repo
+    cluster.createClusterVersion(targetStack, targetRepo, "admin",
+        RepositoryVersionState.INSTALLING);
+    cluster.transitionClusterVersion(targetStack, targetRepo, RepositoryVersionState.INSTALLED);
+
+    cluster.mapHostVersions(Collections.singleton(hostName), cluster.getCurrentClusterVersion(),
+        RepositoryVersionState.CURRENT);
+
+
     Service service = installService(cluster, "HDFS");
     addServiceComponent(cluster, service, "NAMENODE");
     addServiceComponent(cluster, service, "DATANODE");
     createNewServiceComponentHost(cluster, "HDFS", "NAMENODE", "h1");
     createNewServiceComponentHost(cluster, "HDFS", "DATANODE", "h1");
 
-    createNewRepoVersion(targetStack, targetRepo, clusterName, hostName);
-
     // create some configs
     createConfigs(cluster);
+
+    // install the target repo
+    installRepositoryOnHost(hostName, targetRepoVersion);
+
     // setup the cluster for the upgrade across stacks
     cluster.setCurrentStackVersion(sourceStack);
     cluster.setDesiredStackVersion(targetStack);
@@ -348,7 +406,7 @@ public class ComponentVersionCheckActionTest {
 
     // now finalize and ensure we can transition from UPGRADING to UPGRADED
     // automatically before CURRENT
-    Map<String, String> commandParams = new HashMap<String, String>();
+    Map<String, String> commandParams = new HashMap<>();
     commandParams.put(FinalizeUpgradeAction.UPGRADE_DIRECTION_KEY, "upgrade");
     commandParams.put(FinalizeUpgradeAction.VERSION_KEY, targetRepo);
     commandParams.put(FinalizeUpgradeAction.ORIGINAL_STACK_KEY, sourceStack.getStackId());
@@ -417,8 +475,8 @@ public class ComponentVersionCheckActionTest {
   }
 
   private void createConfigs(Cluster cluster) {
-    Map<String, String> properties = new HashMap<String, String>();
-    Map<String, Map<String, String>> propertiesAttributes = new HashMap<String, Map<String, String>>();
+    Map<String, String> properties = new HashMap<>();
+    Map<String, Map<String, String>> propertiesAttributes = new HashMap<>();
     properties.put("a", "a1");
     properties.put("b", "b1");
 

http://git-wip-us.apache.org/repos/asf/ambari/blob/77099797/ambari-server/src/test/java/org/apache/ambari/server/serveraction/upgrades/UpgradeActionTest.java
----------------------------------------------------------------------
diff --git a/ambari-server/src/test/java/org/apache/ambari/server/serveraction/upgrades/UpgradeActionTest.java b/ambari-server/src/test/java/org/apache/ambari/server/serveraction/upgrades/UpgradeActionTest.java
index d875443..46949ea 100644
--- a/ambari-server/src/test/java/org/apache/ambari/server/serveraction/upgrades/UpgradeActionTest.java
+++ b/ambari-server/src/test/java/org/apache/ambari/server/serveraction/upgrades/UpgradeActionTest.java
@@ -70,7 +70,6 @@ import org.apache.ambari.server.state.Clusters;
 import org.apache.ambari.server.state.Config;
 import org.apache.ambari.server.state.ConfigFactory;
 import org.apache.ambari.server.state.Host;
-import org.apache.ambari.server.state.RepositoryInfo;
 import org.apache.ambari.server.state.RepositoryVersionState;
 import org.apache.ambari.server.state.Service;
 import org.apache.ambari.server.state.ServiceComponent;
@@ -81,6 +80,7 @@ import org.apache.ambari.server.state.ServiceFactory;
 import org.apache.ambari.server.state.StackId;
 import org.apache.ambari.server.state.State;
 import org.apache.ambari.server.state.stack.UpgradePack;
+import org.apache.ambari.server.state.stack.upgrade.Direction;
 import org.apache.ambari.server.state.stack.upgrade.UpgradeType;
 import org.apache.ambari.server.utils.EventBusSynchronizer;
 import org.junit.After;
@@ -200,12 +200,10 @@ public class UpgradeActionTest {
     host.setHostAttributes(hostAttributes);
 
     // Create the starting repo version
-    m_helper.getOrCreateRepositoryVersion(sourceStack, sourceRepo);
     c.createClusterVersion(sourceStack, sourceRepo, "admin", RepositoryVersionState.INSTALLING);
     c.transitionClusterVersion(sourceStack, sourceRepo, RepositoryVersionState.CURRENT);
 
     // Start upgrading the newer repo
-    m_helper.getOrCreateRepositoryVersion(targetStack, targetRepo);
     c.createClusterVersion(targetStack, targetRepo, "admin", RepositoryVersionState.INSTALLING);
     c.transitionClusterVersion(targetStack, targetRepo, RepositoryVersionState.INSTALLED);
 
@@ -237,12 +235,10 @@ public class UpgradeActionTest {
     host.setHostAttributes(hostAttributes);
 
     // Create the starting repo version
-    m_helper.getOrCreateRepositoryVersion(sourceStack, sourceRepo);
     c.createClusterVersion(sourceStack, sourceRepo, "admin", RepositoryVersionState.INSTALLING);
     c.transitionClusterVersion(sourceStack, sourceRepo, RepositoryVersionState.CURRENT);
 
     // Start upgrading the mid repo
-    m_helper.getOrCreateRepositoryVersion(midStack, midRepo);
     c.setDesiredStackVersion(midStack);
     c.createClusterVersion(midStack, midRepo, "admin", RepositoryVersionState.INSTALLING);
     c.transitionClusterVersion(midStack, midRepo, RepositoryVersionState.INSTALLED);
@@ -253,8 +249,6 @@ public class UpgradeActionTest {
 
     // Notice that we have not yet changed the cluster current stack to the mid stack to simulate
     // the user skipping this step.
-
-    m_helper.getOrCreateRepositoryVersion(targetStack, targetRepo);
     c.setDesiredStackVersion(targetStack);
     c.createClusterVersion(targetStack, targetRepo, "admin", RepositoryVersionState.INSTALLING);
     c.transitionClusterVersion(targetStack, targetRepo, RepositoryVersionState.INSTALLED);
@@ -301,23 +295,6 @@ public class UpgradeActionTest {
 
     // without this, HostEntity will not have a relation to ClusterEntity
     clusters.mapHostToCluster(hostName, clusterName);
-
-    // Create the starting repo version
-    RepositoryVersionEntity repoEntity = m_helper.getOrCreateRepositoryVersion(sourceStack, sourceRepo);
-    repoEntity.setOperatingSystems("[\n" +
-            "   {\n" +
-            "      \"repositories\":[\n" +
-            "         {\n" +
-            "            \"Repositories/base_url\":\"http://s3.amazonaws.com/dev.hortonworks.com/HDP/centos5/2.x/updates/2.2.0.0\",\n" +
-            "            \"Repositories/repo_name\":\"HDP\",\n" +
-            "            \"Repositories/repo_id\":\"HDP-2.2\"\n" +
-            "         }\n" +
-            "      ],\n" +
-            "      \"OperatingSystems/os_type\":\"redhat6\"\n" +
-            "   }\n" +
-            "]");
-    repoVersionDAO.merge(repoEntity);
-
     c.createClusterVersion(sourceStack, sourceRepo, "admin", RepositoryVersionState.INSTALLING);
     c.transitionClusterVersion(sourceStack, sourceRepo, RepositoryVersionState.CURRENT);
 
@@ -329,13 +306,6 @@ public class UpgradeActionTest {
     StackEntity stackEntityTarget = stackDAO.find(targetStack.getStackName(), targetStack.getStackVersion());
     assertNotNull(stackEntityTarget);
 
-    // Create the new repo version
-    String urlInfo = "[{'repositories':["
-            + "{'Repositories/base_url':'http://foo1','Repositories/repo_name':'HDP','Repositories/repo_id':'" + targetStack.getStackId() + "'}"
-            + "], 'OperatingSystems/os_type':'redhat6'}]";
-
-    repoVersionDAO.create(stackEntityTarget, targetRepo, String.valueOf(System.currentTimeMillis()), urlInfo);
-
     // Start upgrading the newer repo
     c.createClusterVersion(targetStack, targetRepo, "admin", RepositoryVersionState.INSTALLING);
     c.transitionClusterVersion(targetStack, targetRepo, RepositoryVersionState.INSTALLED);
@@ -388,7 +358,6 @@ public class UpgradeActionTest {
     clusters.mapHostToCluster(hostName, clusterName);
 
     // Create the starting repo version
-    m_helper.getOrCreateRepositoryVersion(sourceStack, sourceRepo);
     c.createClusterVersion(sourceStack, sourceRepo, "admin", RepositoryVersionState.INSTALLING);
     c.transitionClusterVersion(sourceStack, sourceRepo, RepositoryVersionState.CURRENT);
   }
@@ -398,13 +367,6 @@ public class UpgradeActionTest {
     assertNotNull(stackEntityTarget);
     Cluster c = clusters.getCluster(clusterName);
 
-
-    // Create the new repo version
-    String urlInfo = "[{'repositories':["
-        + "{'Repositories/base_url':'http://foo1','Repositories/repo_name':'HDP','Repositories/repo_id':'" + targetRepo + "'}"
-        + "], 'OperatingSystems/os_type':'redhat6'}]";
-    repoVersionDAO.create(stackEntityTarget, targetRepo, String.valueOf(System.currentTimeMillis()), urlInfo);
-
     // Start upgrading the newer repo
     c.createClusterVersion(targetStack, targetRepo, "admin", RepositoryVersionState.INSTALLING);
     c.transitionClusterVersion(targetStack, targetRepo, RepositoryVersionState.INSTALLED);
@@ -432,6 +394,9 @@ public class UpgradeActionTest {
     StackId targetStack = HDP_22_STACK;
     String sourceRepo = HDP_2_1_1_0;
     String targetRepo = HDP_2_2_0_1;
+    RepositoryVersionEntity sourceRepositoryVersion = repositoryVersion2110;
+    RepositoryVersionEntity targetRepositoryVersion = repositoryVersion2201;
+
     String hostName = "h1";
 
     // Must be a NON_ROLLING upgrade that jumps stacks in order for it to apply config changes.
@@ -442,7 +407,6 @@ public class UpgradeActionTest {
     Assert.assertTrue(packs.containsKey(upgradePackName));
 
     makeCrossStackUpgradeClusterAndSourceRepo(sourceStack, sourceRepo, hostName);
-//    makeCrossStackUpgradeCluster(sourceStack, sourceRepo, targetStack, targetRepo);
 
     Cluster cluster = clusters.getCluster(clusterName);
 
@@ -460,6 +424,8 @@ public class UpgradeActionTest {
     createNewServiceComponentHost(cluster, "HDFS", "DATANODE", "h1");
 
     makeCrossStackUpgradeTargetRepo(targetStack, targetRepo, hostName);
+    createUpgrade(cluster, upgradePackName, Direction.UPGRADE, sourceRepositoryVersion,
+        targetRepositoryVersion);
 
     RepositoryVersionEntity targetRve = repoVersionDAO.findByStackNameAndVersion("HDP", targetRepo);
     Assert.assertNotNull(targetRve);
@@ -509,8 +475,12 @@ public class UpgradeActionTest {
     StackId targetStack = HDP_21_STACK;
     String sourceRepo = HDP_2_1_1_0;
     String targetRepo = HDP_2_1_1_1;
+    RepositoryVersionEntity sourceRepositoryVersion = repositoryVersion2110;
+    RepositoryVersionEntity targetRepositoryVersion = repositoryVersion2111;
 
     makeDowngradeCluster(sourceStack, sourceRepo, targetStack, targetRepo);
+    Cluster cluster = clusters.getCluster(clusterName);
+    createUpgrade(cluster, sourceRepositoryVersion, targetRepositoryVersion);
 
     Map<String, String> commandParams = new HashMap<>();
     commandParams.put(FinalizeUpgradeAction.UPGRADE_DIRECTION_KEY, "downgrade");
@@ -598,6 +568,9 @@ public class UpgradeActionTest {
     StackId targetStack = HDP_21_STACK;
     String sourceRepo = HDP_2_1_1_0;
     String targetRepo = HDP_2_1_1_1;
+    RepositoryVersionEntity sourceRepositoryVersion = repositoryVersion2110;
+    RepositoryVersionEntity targetRepositoryVersion = repositoryVersion2111;
+
     String hostName = "h1";
 
     createUpgradeClusterAndSourceRepo(sourceStack, sourceRepo, hostName);
@@ -608,9 +581,7 @@ public class UpgradeActionTest {
     Host host = clusters.getHost("h1");
     Cluster cluster = clusters.getCluster(clusterName);
 
-    RepositoryInfo repo = ambariMetaInfo.getRepository(sourceStack.getStackName(), sourceStack.getStackVersion(), "redhat6", sourceStack.getStackId());
-    assertEquals(HDP_211_CENTOS6_REPO_URL, repo.getBaseUrl());
-    verifyBaseRepoURL(helper, cluster, host, HDP_211_CENTOS6_REPO_URL);
+    createUpgrade(cluster, sourceRepositoryVersion, targetRepositoryVersion);
 
     // Finalize the upgrade
     Map<String, String> commandParams = new HashMap<>();
@@ -647,6 +618,8 @@ public class UpgradeActionTest {
     StackId targetStack = HDP_21_STACK;
     String sourceRepo = HDP_2_1_1_0;
     String targetRepo = HDP_2_1_1_1;
+    RepositoryVersionEntity sourceRepositoryVersion = repositoryVersion2110;
+    RepositoryVersionEntity targetRepositoryVersion = repositoryVersion2111;
     String hostName = "h1";
 
     createUpgradeClusterAndSourceRepo(sourceStack, sourceRepo, hostName);
@@ -670,11 +643,7 @@ public class UpgradeActionTest {
     AmbariCustomCommandExecutionHelper helper = m_injector.getInstance(AmbariCustomCommandExecutionHelper.class);
     Host host = clusters.getHost("h1");
     Cluster cluster = clusters.getCluster(clusterName);
-
-    RepositoryInfo repo = ambariMetaInfo.getRepository(sourceStack.getStackName(),
-            sourceStack.getStackVersion(), "redhat6", sourceStack.getStackId());
-    assertEquals(HDP_211_CENTOS6_REPO_URL, repo.getBaseUrl());
-    verifyBaseRepoURL(helper, cluster, host, HDP_211_CENTOS6_REPO_URL);
+    createUpgrade(cluster, sourceRepositoryVersion, targetRepositoryVersion);
 
     // Finalize the upgrade
     Map<String, String> commandParams = new HashMap<>();
@@ -715,12 +684,16 @@ public class UpgradeActionTest {
     StackId targetStack = HDP_22_STACK;
     String sourceRepo = HDP_2_1_1_0;
     String targetRepo = HDP_2_2_0_1;
+    RepositoryVersionEntity sourceRepositoryVersion = repositoryVersion2110;
+    RepositoryVersionEntity targetRepositoryVersion = repositoryVersion2201;
+
     String hostName = "h1";
 
     makeCrossStackUpgradeClusterAndSourceRepo(sourceStack, sourceRepo, hostName);
     makeCrossStackUpgradeTargetRepo(targetStack, targetRepo, hostName);
 
     Cluster cluster = clusters.getCluster(clusterName);
+    createUpgrade(cluster, sourceRepositoryVersion, targetRepositoryVersion);
 
     // setup the cluster for the upgrade across stacks
     cluster.setCurrentStackVersion(sourceStack);
@@ -768,12 +741,14 @@ public class UpgradeActionTest {
     StackId targetStack = HDP_22_STACK;
     String sourceRepo = HDP_2_1_1_0;
     String targetRepo = HDP_2_2_0_1;
+    RepositoryVersionEntity sourceRepositoryVersion = repositoryVersion2110;
+    RepositoryVersionEntity targetRepositoryVersion = repositoryVersion2201;
+
     String hostName = "h1";
 
     makeCrossStackUpgradeClusterAndSourceRepo(sourceStack, sourceRepo, hostName);
     Cluster cluster = clusters.getCluster(clusterName);
 
-
     // install HDFS with some components
     Service service = installService(cluster, "HDFS");
     addServiceComponent(cluster, service, "NAMENODE");
@@ -799,14 +774,10 @@ public class UpgradeActionTest {
     Collection<Config> configs = cluster.getAllConfigs();
     assertEquals(8, configs.size());
 
-    Map<String, String> commandParams = new HashMap<>();
-    commandParams.put(FinalizeUpgradeAction.UPGRADE_DIRECTION_KEY, "downgrade");
-    commandParams.put(FinalizeUpgradeAction.VERSION_KEY, sourceRepo);
-    commandParams.put(FinalizeUpgradeAction.ORIGINAL_STACK_KEY, sourceStack.getStackId());
-    commandParams.put(FinalizeUpgradeAction.TARGET_STACK_KEY, targetStack.getStackId());
+    createUpgrade(cluster, "", Direction.DOWNGRADE, targetRepositoryVersion,
+        sourceRepositoryVersion);
 
     ExecutionCommand executionCommand = new ExecutionCommand();
-    executionCommand.setCommandParams(commandParams);
     executionCommand.setClusterName(clusterName);
 
     HostRoleCommand hostRoleCommand = hostRoleCommandFactory.create(null, null, null, null);
@@ -861,6 +832,9 @@ public class UpgradeActionTest {
     StackId targetStack = HDP_22_STACK;
     String sourceRepo = HDP_2_1_1_0;
     String targetRepo = HDP_2_2_0_1;
+    RepositoryVersionEntity sourceRepositoryVersion = repositoryVersion2110;
+    RepositoryVersionEntity targetRepositoryVersion = repositoryVersion2201;
+
     String hostName = "h1";
 
     makeCrossStackUpgradeClusterAndSourceRepo(sourceStack, sourceRepo, hostName);
@@ -874,6 +848,8 @@ public class UpgradeActionTest {
     createNewServiceComponentHost(cluster, "HDFS", "DATANODE", "h1");
 
     makeCrossStackUpgradeTargetRepo(targetStack, targetRepo, hostName);
+    createUpgrade(cluster, sourceRepositoryVersion, targetRepositoryVersion);
+
     // create some configs
     createConfigs(cluster);
 
@@ -985,6 +961,8 @@ public class UpgradeActionTest {
     upgradeEntity.setUpgradeType(UpgradeType.NON_ROLLING);
     upgradeDAO.create(upgradeEntity);
 
+    cluster.setUpgradeEntity(upgradeEntity);
+
     // verify that no history exist exists yet
     List<ServiceComponentHistoryEntity> historyEntites = serviceComponentDesiredStateDAO.findHistory(
             cluster.getClusterId(), nnSCH.getServiceName(),
@@ -992,10 +970,6 @@ public class UpgradeActionTest {
 
     assertEquals(0, historyEntites.size());
 
-    RepositoryInfo repo = ambariMetaInfo.getRepository(sourceStack.getStackName(), sourceStack.getStackVersion(), "redhat6", sourceStack.getStackId());
-    assertEquals(HDP_211_CENTOS6_REPO_URL, repo.getBaseUrl());
-    verifyBaseRepoURL(helper, cluster, host, HDP_211_CENTOS6_REPO_URL);
-
     // Finalize the upgrade, passing in the request ID so that history is
     // created
     Map<String, String> commandParams = new HashMap<>();
@@ -1017,9 +991,6 @@ public class UpgradeActionTest {
     assertNotNull(report);
     assertEquals(HostRoleStatus.COMPLETED.name(), report.getStatus());
 
-    // Verify the metainfo url
-    verifyBaseRepoURL(helper, cluster, host, "http://foo1");
-
     // ensure that history now exists
     historyEntites = serviceComponentDesiredStateDAO.findHistory(cluster.getClusterId(),
             nnSCH.getServiceName(), nnSCH.getServiceComponentName());
@@ -1095,4 +1066,45 @@ public class UpgradeActionTest {
     configFactory.createNew(cluster, "foo-site", "version-" + System.currentTimeMillis(),
         properties, propertiesAttributes);
   }
+
+  /**
+   * Creates an upgrade and associates it with the cluster.
+   */
+  private UpgradeEntity createUpgrade(Cluster cluster,
+      RepositoryVersionEntity sourceRepositoryVersion,
+      RepositoryVersionEntity targetRepositoryVersion) throws Exception {
+    return createUpgrade(cluster, "", Direction.UPGRADE, sourceRepositoryVersion,
+        targetRepositoryVersion);
+  }
+
+  /**
+   * Creates an upgrade and associates it with the cluster.
+   */
+  private UpgradeEntity createUpgrade(Cluster cluster, String upgradePackName, Direction direction,
+      RepositoryVersionEntity sourceRepositoryVersion,
+      RepositoryVersionEntity targetRepositoryVersion) throws Exception {
+
+    // create some entities for the finalize action to work with for patch
+    // history
+    RequestEntity requestEntity = new RequestEntity();
+    requestEntity.setClusterId(cluster.getClusterId());
+    requestEntity.setRequestId(1L);
+    requestEntity.setStartTime(System.currentTimeMillis());
+    requestEntity.setCreateTime(System.currentTimeMillis());
+    requestDAO.create(requestEntity);
+
+    UpgradeEntity upgradeEntity = new UpgradeEntity();
+    upgradeEntity.setId(1L);
+    upgradeEntity.setClusterId(cluster.getClusterId());
+    upgradeEntity.setRequestEntity(requestEntity);
+    upgradeEntity.setUpgradePackage(upgradePackName);
+    upgradeEntity.setFromRepositoryVersion(sourceRepositoryVersion);
+    upgradeEntity.setToRepositoryVersion(targetRepositoryVersion);
+    upgradeEntity.setUpgradeType(UpgradeType.NON_ROLLING);
+    upgradeEntity.setDirection(direction);
+
+    upgradeDAO.create(upgradeEntity);
+    cluster.setUpgradeEntity(upgradeEntity);
+    return upgradeEntity;
+  }
 }

http://git-wip-us.apache.org/repos/asf/ambari/blob/77099797/ambari-server/src/test/java/org/apache/ambari/server/state/UpgradeHelperTest.java
----------------------------------------------------------------------
diff --git a/ambari-server/src/test/java/org/apache/ambari/server/state/UpgradeHelperTest.java b/ambari-server/src/test/java/org/apache/ambari/server/state/UpgradeHelperTest.java
index b5dec7f..8b1f4ef 100644
--- a/ambari-server/src/test/java/org/apache/ambari/server/state/UpgradeHelperTest.java
+++ b/ambari-server/src/test/java/org/apache/ambari/server/state/UpgradeHelperTest.java
@@ -18,6 +18,7 @@
 package org.apache.ambari.server.state;
 
 import static org.easymock.EasyMock.expect;
+import static org.easymock.EasyMock.expectLastCall;
 import static org.easymock.EasyMock.replay;
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertFalse;
@@ -40,6 +41,7 @@ import java.util.regex.Pattern;
 
 import org.apache.ambari.server.AmbariException;
 import org.apache.ambari.server.H2DatabaseCleaner;
+import org.apache.ambari.server.actionmanager.HostRoleCommandFactory;
 import org.apache.ambari.server.api.services.AmbariMetaInfo;
 import org.apache.ambari.server.controller.AmbariManagementController;
 import org.apache.ambari.server.controller.ClusterRequest;
@@ -53,6 +55,7 @@ import org.apache.ambari.server.security.authorization.AuthorizationException;
 import org.apache.ambari.server.stack.HostsType;
 import org.apache.ambari.server.stack.MasterHostResolver;
 import org.apache.ambari.server.stack.StackManagerMock;
+import org.apache.ambari.server.stageplanner.RoleGraphFactory;
 import org.apache.ambari.server.state.UpgradeHelper.UpgradeGroupHolder;
 import org.apache.ambari.server.state.stack.ConfigUpgradePack;
 import org.apache.ambari.server.state.stack.UpgradePack;
@@ -74,7 +77,9 @@ import org.apache.ambari.server.state.stack.upgrade.TaskWrapper;
 import org.apache.ambari.server.state.stack.upgrade.UpgradeScope;
 import org.apache.ambari.server.state.stack.upgrade.UpgradeType;
 import org.apache.ambari.server.utils.EventBusSynchronizer;
+import org.easymock.Capture;
 import org.easymock.EasyMock;
+import org.easymock.IAnswer;
 import org.junit.After;
 import org.junit.Before;
 import org.junit.Test;
@@ -291,12 +296,8 @@ public class UpgradeHelperTest {
     Cluster cluster = makeCluster();
 
     UpgradeContext context = getMockUpgradeContext(cluster, Direction.UPGRADE, UpgradeType.ROLLING,
-        repositoryVersion2210, m_masterHostResolver, false);
-
-    context.setSupportedServices(Collections.singleton("ZOOKEEPER"));
-    context.setScope(UpgradeScope.PARTIAL);
-
-    replay(context);
+        repositoryVersion2210, Collections.singleton("ZOOKEEPER"), m_masterHostResolver,
+        UpgradeScope.PARTIAL, true);
 
     List<Grouping> groupings = upgrade.getGroups(Direction.UPGRADE);
     assertEquals(8, groupings.size());
@@ -351,12 +352,9 @@ public class UpgradeHelperTest {
     Cluster cluster = makeCluster();
 
     UpgradeContext context = getMockUpgradeContext(cluster, Direction.UPGRADE, UpgradeType.ROLLING,
-        repositoryVersion2210, m_masterHostResolver, false);
+        repositoryVersion2210, Collections.singleton("ZOOKEEPER"), m_masterHostResolver,
+        UpgradeScope.COMPLETE, true);
 
-    context.setSupportedServices(Collections.singleton("ZOOKEEPER"));
-    context.setScope(UpgradeScope.COMPLETE);
-
-    replay(context);
 
     List<Grouping> groupings = upgrade.getGroups(Direction.UPGRADE);
     assertEquals(8, groupings.size());
@@ -1739,11 +1737,9 @@ public class UpgradeHelperTest {
 
     // get an upgrade
     UpgradeContext context = getMockUpgradeContext(cluster, Direction.UPGRADE, UpgradeType.ROLLING,
-        repositoryVersion2210, m_masterHostResolver, false);
+        repositoryVersion2210, Collections.singleton("ZOOKEEPER"), m_masterHostResolver,
+        UpgradeScope.COMPLETE, true);
 
-    context.setSupportedServices(Collections.singleton("ZOOKEEPER"));
-    context.setScope(UpgradeScope.COMPLETE);
-    replay(context);
 
     List<Grouping> groupings = upgradePack.getGroups(Direction.UPGRADE);
     assertEquals(2, groupings.size());
@@ -2164,7 +2160,6 @@ public class UpgradeHelperTest {
     context = getMockUpgradeContext(c, Direction.DOWNGRADE, UpgradeType.HOST_ORDERED,
         repositoryVersion2110, resolver, true);
 
-    context.setResolver(resolver);
     groups = m_upgradeHelper.createSequence(upgradePack, context);
 
     assertEquals(1, groups.size());
@@ -2241,20 +2236,23 @@ public class UpgradeHelperTest {
    * @return
    */
   private UpgradeContext getMockUpgradeContext(Cluster cluster, Direction direction,
-      UpgradeType type) {
-    return getMockUpgradeContext(cluster, direction, type, repositoryVersion2210);
+      UpgradeType type, RepositoryVersionEntity repositoryVersion) {
+    return getMockUpgradeContext(cluster, direction, type, repositoryVersion, m_masterHostResolver, true);
   }
 
   /**
    * @param cluster
    * @param direction
    * @param type
+   * @param repositoryType
+   * @param services
    * @return
    */
   private UpgradeContext getMockUpgradeContext(Cluster cluster, Direction direction,
-      UpgradeType type, RepositoryVersionEntity repositoryVersion) {
-    return getMockUpgradeContext(cluster, direction, type, repositoryVersion, m_masterHostResolver,
-        true);
+      UpgradeType type, RepositoryVersionEntity repositoryVersion, MasterHostResolver resolver,
+      boolean replay) {
+    return getMockUpgradeContext(cluster, direction, type, repositoryVersion,
+        new HashSet<String>(), resolver, UpgradeScope.COMPLETE, true);
   }
 
   /**
@@ -2266,8 +2264,8 @@ public class UpgradeHelperTest {
    * @return
    */
   private UpgradeContext getMockUpgradeContext(Cluster cluster, Direction direction,
-      UpgradeType type, RepositoryVersionEntity repositoryVersion, MasterHostResolver resolver,
-      boolean replay) {
+      UpgradeType type, RepositoryVersionEntity repositoryVersion, final Set<String> services,
+      MasterHostResolver resolver, final UpgradeScope scope, boolean replay) {
     RepositoryVersionEntity sourceRepositoryVersion = cluster.getAllClusterVersions().iterator().next().getRepositoryVersion();
 
     UpgradeContext context = EasyMock.createNiceMock(UpgradeContext.class);
@@ -2276,13 +2274,96 @@ public class UpgradeHelperTest {
     expect(context.getDirection()).andReturn(direction).anyTimes();
     expect(context.getSourceRepositoryVersion()).andReturn(sourceRepositoryVersion).anyTimes();
     expect(context.getTargetRepositoryVersion()).andReturn(repositoryVersion).anyTimes();
+    expect(context.getEffectiveStackId()).andReturn(repositoryVersion.getStackId()).anyTimes();
+    expect(context.getTargetStackId()).andReturn(repositoryVersion.getStackId()).anyTimes();
     expect(context.getAmbariMetaInfo()).andReturn(ambariMetaInfo).anyTimes();
+    expect(context.getHostRoleCommandFactory()).andStubReturn(injector.getInstance(HostRoleCommandFactory.class));
+    expect(context.getRoleGraphFactory()).andStubReturn(injector.getInstance(RoleGraphFactory.class));
 
     // only set this if supplied
     if (null != resolver) {
       expect(context.getResolver()).andReturn(resolver).anyTimes();
     }
 
+    final Capture<UpgradeScope> isScopedCapture = EasyMock.newCapture();
+    expect(context.isScoped(EasyMock.capture(isScopedCapture))).andStubAnswer(
+        new IAnswer<Boolean>() {
+          @Override
+          public Boolean answer() throws Throwable {
+            UpgradeScope capturedScope = isScopedCapture.getValue();
+            if (capturedScope == UpgradeScope.ANY) {
+              return true;
+            }
+
+            return scope.isScoped(capturedScope);
+          }
+        });
+
+    final Capture<String> serviceNameSupported = EasyMock.newCapture();
+    expect(context.isServiceSupported(EasyMock.capture(serviceNameSupported))).andAnswer(
+        new IAnswer<Boolean>() {
+          @Override
+          public Boolean answer() {
+            if (services.isEmpty()) {
+              return true;
+            }
+
+            return services.contains(serviceNameSupported.getValue());
+          }
+        }).anyTimes();
+
+    final Map<String, String> serviceNames = new HashMap<>();
+    final Capture<String> serviceDisplayNameArg1 = EasyMock.newCapture();
+    final Capture<String> serviceDisplayNameArg2 = EasyMock.newCapture();
+
+    context.setServiceDisplay(EasyMock.capture(serviceDisplayNameArg1),
+        EasyMock.capture(serviceDisplayNameArg2));
+    expectLastCall().andAnswer(new IAnswer<Object>() {
+      @Override
+      public Object answer() {
+        serviceNames.put(serviceDisplayNameArg1.getValue(), serviceDisplayNameArg2.getValue());
+        return null;
+      }
+    }).anyTimes();
+
+    final Map<String, String> componentNames = new HashMap<>();
+    final Capture<String> componentDisplayNameArg1 = EasyMock.newCapture();
+    final Capture<String> componentDisplayNameArg2 = EasyMock.newCapture();
+    final Capture<String> componentDisplayNameArg3 = EasyMock.newCapture();
+
+    context.setComponentDisplay(EasyMock.capture(componentDisplayNameArg1),
+        EasyMock.capture(componentDisplayNameArg2), EasyMock.capture(componentDisplayNameArg3));
+
+    expectLastCall().andAnswer(new IAnswer<Object>() {
+      @Override
+      public Object answer() {
+        componentNames.put(
+            componentDisplayNameArg1.getValue() + ":" + componentDisplayNameArg2.getValue(),
+            componentDisplayNameArg3.getValue());
+        return null;
+      }
+    }).anyTimes();
+
+    final Capture<String> getServiceDisplayArgument1 = EasyMock.newCapture();
+    expect(context.getServiceDisplay(EasyMock.capture(getServiceDisplayArgument1))).andAnswer(
+        new IAnswer<String>() {
+          @Override
+          public String answer() {
+            return serviceNames.get(getServiceDisplayArgument1.getValue());
+          }
+        }).anyTimes();
+
+    final Capture<String> getComponentDisplayArgument1 = EasyMock.newCapture();
+    final Capture<String> getComponentDisplayArgument2 = EasyMock.newCapture();
+    expect(context.getComponentDisplay(EasyMock.capture(getComponentDisplayArgument1),
+        EasyMock.capture(getComponentDisplayArgument2))).andAnswer(new IAnswer<String>() {
+          @Override
+          public String answer() {
+            return componentNames.get(getComponentDisplayArgument1.getValue() + ":"
+                + getComponentDisplayArgument2.getValue());
+          }
+        }).anyTimes();
+
     if (replay) {
       replay(context);
     }

http://git-wip-us.apache.org/repos/asf/ambari/blob/77099797/ambari-server/src/test/java/org/apache/ambari/server/state/cluster/ClusterEffectiveVersionTest.java
----------------------------------------------------------------------
diff --git a/ambari-server/src/test/java/org/apache/ambari/server/state/cluster/ClusterEffectiveVersionTest.java b/ambari-server/src/test/java/org/apache/ambari/server/state/cluster/ClusterEffectiveVersionTest.java
index 22183ba..9bb9b45 100644
--- a/ambari-server/src/test/java/org/apache/ambari/server/state/cluster/ClusterEffectiveVersionTest.java
+++ b/ambari-server/src/test/java/org/apache/ambari/server/state/cluster/ClusterEffectiveVersionTest.java
@@ -199,7 +199,7 @@ public class ClusterEffectiveVersionTest extends EasyMockSupport {
 
     EasyMock.expect(upgradeEntity.getFromRepositoryVersion()).andReturn(repoVersion2400).anyTimes();
     EasyMock.expect(upgradeEntity.getToRepositoryVersion()).andReturn(repoVersion2300).atLeastOnce();
-    EasyMock.expect(upgradeEntity.getDirection()).andReturn(Direction.DOWNGRADE).atLeastOnce();
+    EasyMock.expect(upgradeEntity.getDirection()).andReturn(Direction.DOWNGRADE).anyTimes();
 
     ClusterVersionEntity clusterVersionUpgradingTo = createNiceMock(ClusterVersionEntity.class);
     EasyMock.expect(clusterVersionUpgradingTo.getRepositoryVersion()).andReturn(