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/27 23:38:19 UTC

[2/3] ambari git commit: AMBARI-21358 - Upgrades Must Track To/From Repository Versions Instead of Strings (jonathanhurley)

http://git-wip-us.apache.org/repos/asf/ambari/blob/7ad307c2/ambari-server/src/main/java/org/apache/ambari/server/upgrade/UpgradeCatalog252.java
----------------------------------------------------------------------
diff --git a/ambari-server/src/main/java/org/apache/ambari/server/upgrade/UpgradeCatalog252.java b/ambari-server/src/main/java/org/apache/ambari/server/upgrade/UpgradeCatalog252.java
index c25dec1..fd26504 100644
--- a/ambari-server/src/main/java/org/apache/ambari/server/upgrade/UpgradeCatalog252.java
+++ b/ambari-server/src/main/java/org/apache/ambari/server/upgrade/UpgradeCatalog252.java
@@ -33,6 +33,10 @@ public class UpgradeCatalog252 extends AbstractUpgradeCatalog {
   static final String CLUSTERCONFIG_TABLE = "clusterconfig";
   static final String SERVICE_DELETED_COLUMN = "service_deleted";
 
+  private static final String UPGRADE_TABLE = "upgrade";
+  private static final String UPGRADE_TABLE_FROM_REPO_COLUMN = "from_repo_version_id";
+  private static final String UPGRADE_TABLE_TO_REPO_COLUMN = "to_repo_version_id";
+
   /**
    * Constructor.
    *
@@ -65,6 +69,7 @@ public class UpgradeCatalog252 extends AbstractUpgradeCatalog {
   @Override
   protected void executeDDLUpdates() throws AmbariException, SQLException {
     addServiceDeletedColumnToClusterConfigTable();
+    addRepositoryColumnsToUpgradeTable();
   }
 
   /**
@@ -91,4 +96,32 @@ public class UpgradeCatalog252 extends AbstractUpgradeCatalog {
     dbAccessor.addColumn(CLUSTERCONFIG_TABLE,
         new DBColumnInfo(SERVICE_DELETED_COLUMN, Short.class, null, 0, false));
   }
+
+  /**
+   * Changes the following columns to {@value #UPGRADE_TABLE}:
+   * <ul>
+   * <li>{@value #UPGRADE_TABLE_FROM_REPO_COLUMN}
+   * <li>{@value #UPGRADE_TABLE_TO_REPO_COLUMN}
+   * <li>Removes {@code to_version}
+   * <li>Removes {@code from_version}
+   * </ul>
+   *
+   * @throws SQLException
+   */
+  private void addRepositoryColumnsToUpgradeTable() throws SQLException {
+    dbAccessor.dropColumn(UPGRADE_TABLE, "to_version");
+    dbAccessor.dropColumn(UPGRADE_TABLE, "from_version");
+
+    dbAccessor.addColumn(UPGRADE_TABLE,
+        new DBColumnInfo(UPGRADE_TABLE_FROM_REPO_COLUMN, Long.class, null, null, false));
+
+    dbAccessor.addFKConstraint(UPGRADE_TABLE, "FK_upgrade_from_repo_id",
+        UPGRADE_TABLE_FROM_REPO_COLUMN, "repo_version", "repo_version_id", false);
+
+    dbAccessor.addColumn(UPGRADE_TABLE,
+        new DBColumnInfo(UPGRADE_TABLE_TO_REPO_COLUMN, Long.class, null, null, false));
+
+    dbAccessor.addFKConstraint(UPGRADE_TABLE, "FK_upgrade_to_repo_id",
+        UPGRADE_TABLE_FROM_REPO_COLUMN, "repo_version", "repo_version_id", false);
+  }
 }

http://git-wip-us.apache.org/repos/asf/ambari/blob/7ad307c2/ambari-server/src/main/resources/Ambari-DDL-Derby-CREATE.sql
----------------------------------------------------------------------
diff --git a/ambari-server/src/main/resources/Ambari-DDL-Derby-CREATE.sql b/ambari-server/src/main/resources/Ambari-DDL-Derby-CREATE.sql
index c99326e..54e7039 100644
--- a/ambari-server/src/main/resources/Ambari-DDL-Derby-CREATE.sql
+++ b/ambari-server/src/main/resources/Ambari-DDL-Derby-CREATE.sql
@@ -831,8 +831,8 @@ CREATE TABLE upgrade (
   upgrade_id BIGINT NOT NULL,
   cluster_id BIGINT NOT NULL,
   request_id BIGINT NOT NULL,
-  from_version VARCHAR(255) DEFAULT '' NOT NULL,
-  to_version VARCHAR(255) DEFAULT '' NOT NULL,
+  from_repo_version_id BIGINT NOT NULL,
+  to_repo_version_id BIGINT NOT NULL,
   direction VARCHAR(255) DEFAULT 'UPGRADE' NOT NULL,
   upgrade_package VARCHAR(255) NOT NULL,
   upgrade_type VARCHAR(32) NOT NULL,
@@ -842,7 +842,9 @@ CREATE TABLE upgrade (
   suspended SMALLINT DEFAULT 0 NOT NULL,
   CONSTRAINT PK_upgrade PRIMARY KEY (upgrade_id),
   FOREIGN KEY (cluster_id) REFERENCES clusters(cluster_id),
-  FOREIGN KEY (request_id) REFERENCES request(request_id)
+  FOREIGN KEY (request_id) REFERENCES request(request_id),
+  CONSTRAINT FK_upgrade_from_repo_id FOREIGN KEY (from_repo_version_id) REFERENCES repo_version (repo_version_id),
+  CONSTRAINT FK_upgrade_to_repo_id FOREIGN KEY (to_repo_version_id) REFERENCES repo_version (repo_version_id)
 );
 
 CREATE TABLE upgrade_group (

http://git-wip-us.apache.org/repos/asf/ambari/blob/7ad307c2/ambari-server/src/main/resources/Ambari-DDL-MySQL-CREATE.sql
----------------------------------------------------------------------
diff --git a/ambari-server/src/main/resources/Ambari-DDL-MySQL-CREATE.sql b/ambari-server/src/main/resources/Ambari-DDL-MySQL-CREATE.sql
index be0bdd8..d6ac15b 100644
--- a/ambari-server/src/main/resources/Ambari-DDL-MySQL-CREATE.sql
+++ b/ambari-server/src/main/resources/Ambari-DDL-MySQL-CREATE.sql
@@ -849,8 +849,8 @@ CREATE TABLE upgrade (
   upgrade_id BIGINT NOT NULL,
   cluster_id BIGINT NOT NULL,
   request_id BIGINT NOT NULL,
-  from_version VARCHAR(255) DEFAULT '' NOT NULL,
-  to_version VARCHAR(255) DEFAULT '' NOT NULL,
+  from_repo_version_id BIGINT NOT NULL,
+  to_repo_version_id BIGINT NOT NULL,
   direction VARCHAR(255) DEFAULT 'UPGRADE' NOT NULL,
   upgrade_package VARCHAR(255) NOT NULL,
   upgrade_type VARCHAR(32) NOT NULL,
@@ -860,7 +860,9 @@ CREATE TABLE upgrade (
   suspended TINYINT(1) DEFAULT 0 NOT NULL,
   CONSTRAINT PK_upgrade PRIMARY KEY (upgrade_id),
   FOREIGN KEY (cluster_id) REFERENCES clusters(cluster_id),
-  FOREIGN KEY (request_id) REFERENCES request(request_id)
+  FOREIGN KEY (request_id) REFERENCES request(request_id),
+  CONSTRAINT FK_upgrade_from_repo_id FOREIGN KEY (from_repo_version_id) REFERENCES repo_version (repo_version_id),
+  CONSTRAINT FK_upgrade_to_repo_id FOREIGN KEY (to_repo_version_id) REFERENCES repo_version (repo_version_id)
 );
 
 CREATE TABLE upgrade_group (

http://git-wip-us.apache.org/repos/asf/ambari/blob/7ad307c2/ambari-server/src/main/resources/Ambari-DDL-Oracle-CREATE.sql
----------------------------------------------------------------------
diff --git a/ambari-server/src/main/resources/Ambari-DDL-Oracle-CREATE.sql b/ambari-server/src/main/resources/Ambari-DDL-Oracle-CREATE.sql
index b0f0d9f..c107a65 100644
--- a/ambari-server/src/main/resources/Ambari-DDL-Oracle-CREATE.sql
+++ b/ambari-server/src/main/resources/Ambari-DDL-Oracle-CREATE.sql
@@ -828,8 +828,8 @@ CREATE TABLE upgrade (
   upgrade_id NUMBER(19) NOT NULL,
   cluster_id NUMBER(19) NOT NULL,
   request_id NUMBER(19) NOT NULL,
-  from_version VARCHAR2(255) DEFAULT '' NOT NULL,
-  to_version VARCHAR2(255) DEFAULT '' NOT NULL,
+  from_repo_version_id NUMBER(19) NOT NULL,
+  to_repo_version_id NUMBER(19) NOT NULL,
   direction VARCHAR2(255) DEFAULT 'UPGRADE' NOT NULL,
   upgrade_package VARCHAR2(255) NOT NULL,
   upgrade_type VARCHAR2(32) NOT NULL,
@@ -839,7 +839,9 @@ CREATE TABLE upgrade (
   suspended NUMBER(1) DEFAULT 0 NOT NULL,
   CONSTRAINT PK_upgrade PRIMARY KEY (upgrade_id),
   FOREIGN KEY (cluster_id) REFERENCES clusters(cluster_id),
-  FOREIGN KEY (request_id) REFERENCES request(request_id)
+  FOREIGN KEY (request_id) REFERENCES request(request_id),
+  CONSTRAINT FK_upgrade_from_repo_id FOREIGN KEY (from_repo_version_id) REFERENCES repo_version (repo_version_id),
+  CONSTRAINT FK_upgrade_to_repo_id FOREIGN KEY (to_repo_version_id) REFERENCES repo_version (repo_version_id)
 );
 
 CREATE TABLE upgrade_group (

http://git-wip-us.apache.org/repos/asf/ambari/blob/7ad307c2/ambari-server/src/main/resources/Ambari-DDL-Postgres-CREATE.sql
----------------------------------------------------------------------
diff --git a/ambari-server/src/main/resources/Ambari-DDL-Postgres-CREATE.sql b/ambari-server/src/main/resources/Ambari-DDL-Postgres-CREATE.sql
index c680a66..5baf964 100644
--- a/ambari-server/src/main/resources/Ambari-DDL-Postgres-CREATE.sql
+++ b/ambari-server/src/main/resources/Ambari-DDL-Postgres-CREATE.sql
@@ -830,8 +830,8 @@ CREATE TABLE upgrade (
   upgrade_id BIGINT NOT NULL,
   cluster_id BIGINT NOT NULL,
   request_id BIGINT NOT NULL,
-  from_version VARCHAR(255) DEFAULT '' NOT NULL,
-  to_version VARCHAR(255) DEFAULT '' NOT NULL,
+  from_repo_version_id BIGINT NOT NULL,
+  to_repo_version_id BIGINT NOT NULL,
   direction VARCHAR(255) DEFAULT 'UPGRADE' NOT NULL,
   upgrade_package VARCHAR(255) NOT NULL,
   upgrade_type VARCHAR(32) NOT NULL,
@@ -841,7 +841,9 @@ CREATE TABLE upgrade (
   suspended SMALLINT DEFAULT 0 NOT NULL,
   CONSTRAINT PK_upgrade PRIMARY KEY (upgrade_id),
   FOREIGN KEY (cluster_id) REFERENCES clusters(cluster_id),
-  FOREIGN KEY (request_id) REFERENCES request(request_id)
+  FOREIGN KEY (request_id) REFERENCES request(request_id),
+  CONSTRAINT FK_upgrade_from_repo_id FOREIGN KEY (from_repo_version_id) REFERENCES repo_version (repo_version_id),
+  CONSTRAINT FK_upgrade_to_repo_id FOREIGN KEY (to_repo_version_id) REFERENCES repo_version (repo_version_id)
 );
 
 CREATE TABLE upgrade_group (

http://git-wip-us.apache.org/repos/asf/ambari/blob/7ad307c2/ambari-server/src/main/resources/Ambari-DDL-SQLAnywhere-CREATE.sql
----------------------------------------------------------------------
diff --git a/ambari-server/src/main/resources/Ambari-DDL-SQLAnywhere-CREATE.sql b/ambari-server/src/main/resources/Ambari-DDL-SQLAnywhere-CREATE.sql
index c6524b0..dfcd2f2 100644
--- a/ambari-server/src/main/resources/Ambari-DDL-SQLAnywhere-CREATE.sql
+++ b/ambari-server/src/main/resources/Ambari-DDL-SQLAnywhere-CREATE.sql
@@ -827,8 +827,8 @@ CREATE TABLE upgrade (
   upgrade_id NUMERIC(19) NOT NULL,
   cluster_id NUMERIC(19) NOT NULL,
   request_id NUMERIC(19) NOT NULL,
-  from_version VARCHAR(255) DEFAULT '' NOT NULL,
-  to_version VARCHAR(255) DEFAULT '' NOT NULL,
+  from_repo_version_id NUMERIC(19) NOT NULL,
+  to_repo_version_id NUMERIC(19) NOT NULL,
   direction VARCHAR(255) DEFAULT 'UPGRADE' NOT NULL,
   upgrade_type VARCHAR(32) NOT NULL,
   upgrade_package VARCHAR(255) NOT NULL,
@@ -838,7 +838,9 @@ CREATE TABLE upgrade (
   suspended BIT DEFAULT 0 NOT NULL,
   CONSTRAINT PK_upgrade PRIMARY KEY (upgrade_id),
   FOREIGN KEY (cluster_id) REFERENCES clusters(cluster_id),
-  FOREIGN KEY (request_id) REFERENCES request(request_id)
+  FOREIGN KEY (request_id) REFERENCES request(request_id),
+  CONSTRAINT FK_upgrade_from_repo_id FOREIGN KEY (from_repo_version_id) REFERENCES repo_version (repo_version_id),
+  CONSTRAINT FK_upgrade_to_repo_id FOREIGN KEY (to_repo_version_id) REFERENCES repo_version (repo_version_id)
 );
 
 CREATE TABLE upgrade_group (

http://git-wip-us.apache.org/repos/asf/ambari/blob/7ad307c2/ambari-server/src/main/resources/Ambari-DDL-SQLServer-CREATE.sql
----------------------------------------------------------------------
diff --git a/ambari-server/src/main/resources/Ambari-DDL-SQLServer-CREATE.sql b/ambari-server/src/main/resources/Ambari-DDL-SQLServer-CREATE.sql
index 6483645..63c17a9 100644
--- a/ambari-server/src/main/resources/Ambari-DDL-SQLServer-CREATE.sql
+++ b/ambari-server/src/main/resources/Ambari-DDL-SQLServer-CREATE.sql
@@ -849,8 +849,8 @@ CREATE TABLE upgrade (
   upgrade_id BIGINT NOT NULL,
   cluster_id BIGINT NOT NULL,
   request_id BIGINT NOT NULL,
-  from_version VARCHAR(255) DEFAULT '' NOT NULL,
-  to_version VARCHAR(255) DEFAULT '' NOT NULL,
+  from_repo_version_id BIGINT NOT NULL,
+  to_repo_version_id BIGINT NOT NULL,
   direction VARCHAR(255) DEFAULT 'UPGRADE' NOT NULL,
   upgrade_package VARCHAR(255) NOT NULL,
   upgrade_type VARCHAR(32) NOT NULL,
@@ -860,7 +860,9 @@ CREATE TABLE upgrade (
   suspended BIT DEFAULT 0 NOT NULL,
   CONSTRAINT PK_upgrade PRIMARY KEY CLUSTERED (upgrade_id),
   FOREIGN KEY (cluster_id) REFERENCES clusters(cluster_id),
-  FOREIGN KEY (request_id) REFERENCES request(request_id)
+  FOREIGN KEY (request_id) REFERENCES request(request_id),
+  CONSTRAINT FK_upgrade_from_repo_id FOREIGN KEY (from_repo_version_id) REFERENCES repo_version (repo_version_id),
+  CONSTRAINT FK_upgrade_to_repo_id FOREIGN KEY (to_repo_version_id) REFERENCES repo_version (repo_version_id)
 );
 
 CREATE TABLE upgrade_group (

http://git-wip-us.apache.org/repos/asf/ambari/blob/7ad307c2/ambari-server/src/test/java/org/apache/ambari/server/alerts/ComponentVersionAlertRunnableTest.java
----------------------------------------------------------------------
diff --git a/ambari-server/src/test/java/org/apache/ambari/server/alerts/ComponentVersionAlertRunnableTest.java b/ambari-server/src/test/java/org/apache/ambari/server/alerts/ComponentVersionAlertRunnableTest.java
index 98f6f44..6c07138 100644
--- a/ambari-server/src/test/java/org/apache/ambari/server/alerts/ComponentVersionAlertRunnableTest.java
+++ b/ambari-server/src/test/java/org/apache/ambari/server/alerts/ComponentVersionAlertRunnableTest.java
@@ -224,8 +224,9 @@ public class ComponentVersionAlertRunnableTest extends EasyMockSupport {
   @Test
   public void testUpgradeInProgress() throws Exception {
     UpgradeEntity upgrade = createNiceMock(UpgradeEntity.class);
-    expect(upgrade.getToVersion()).andReturn("VERSION").once();
-    expect(m_cluster.getUpgradeInProgress()).andReturn(upgrade).once();
+    expect(upgrade.getToRepositoryVersion()).andReturn(EasyMock.createNiceMock(RepositoryVersionEntity.class)).atLeastOnce();
+    expect(upgrade.getFromRepositoryVersion()).andReturn(EasyMock.createNiceMock(RepositoryVersionEntity.class)).atLeastOnce();
+    expect(m_cluster.getUpgradeInProgress()).andReturn(upgrade).atLeastOnce();
 
     replayAll();
 

http://git-wip-us.apache.org/repos/asf/ambari/blob/7ad307c2/ambari-server/src/test/java/org/apache/ambari/server/audit/request/creator/UpgradeEventCreatorTest.java
----------------------------------------------------------------------
diff --git a/ambari-server/src/test/java/org/apache/ambari/server/audit/request/creator/UpgradeEventCreatorTest.java b/ambari-server/src/test/java/org/apache/ambari/server/audit/request/creator/UpgradeEventCreatorTest.java
index 4822584..fdc4752 100644
--- a/ambari-server/src/test/java/org/apache/ambari/server/audit/request/creator/UpgradeEventCreatorTest.java
+++ b/ambari-server/src/test/java/org/apache/ambari/server/audit/request/creator/UpgradeEventCreatorTest.java
@@ -18,8 +18,6 @@
 
 package org.apache.ambari.server.audit.request.creator;
 
-import junit.framework.Assert;
-
 import java.util.HashMap;
 import java.util.Map;
 
@@ -33,6 +31,8 @@ import org.apache.ambari.server.controller.internal.UpgradeResourceProvider;
 import org.apache.ambari.server.controller.spi.Resource;
 import org.junit.Test;
 
+import junit.framework.Assert;
+
 public class UpgradeEventCreatorTest extends AuditEventCreatorTestBase{
 
   @Test

http://git-wip-us.apache.org/repos/asf/ambari/blob/7ad307c2/ambari-server/src/test/java/org/apache/ambari/server/checks/PreviousUpgradeCompletedTest.java
----------------------------------------------------------------------
diff --git a/ambari-server/src/test/java/org/apache/ambari/server/checks/PreviousUpgradeCompletedTest.java b/ambari-server/src/test/java/org/apache/ambari/server/checks/PreviousUpgradeCompletedTest.java
index 4bfa8d4..21b2514 100644
--- a/ambari-server/src/test/java/org/apache/ambari/server/checks/PreviousUpgradeCompletedTest.java
+++ b/ambari-server/src/test/java/org/apache/ambari/server/checks/PreviousUpgradeCompletedTest.java
@@ -18,6 +18,7 @@
 package org.apache.ambari.server.checks;
 
 import org.apache.ambari.server.controller.PrereqCheckRequest;
+import org.apache.ambari.server.orm.entities.RepositoryVersionEntity;
 import org.apache.ambari.server.orm.entities.StackEntity;
 import org.apache.ambari.server.orm.entities.UpgradeEntity;
 import org.apache.ambari.server.state.Cluster;
@@ -43,8 +44,10 @@ public class PreviousUpgradeCompletedTest {
   private final Cluster cluster = Mockito.mock(Cluster.class);
   private StackId sourceStackId = new StackId("HDP", "2.2");
   private StackId targetStackId = new StackId("HDP", "2.2");
-  private String sourceRepositoryVersion = "2.2.6.0-1234";
-  private String destRepositoryVersion = "2.2.8.0-5678";
+  private final RepositoryVersionEntity sourceRepositoryVersion = Mockito.mock(RepositoryVersionEntity.class);
+  private final RepositoryVersionEntity targetRepositoryVersion = Mockito.mock(RepositoryVersionEntity.class);
+  private String sourceVersion = "2.2.6.0-1234";
+  private String destVersion = "2.2.8.0-5678";
   private String clusterName = "cluster";
   private PrereqCheckRequest checkRequest = new PrereqCheckRequest(clusterName);
   private PreviousUpgradeCompleted puc = new PreviousUpgradeCompleted();
@@ -64,7 +67,7 @@ public class PreviousUpgradeCompletedTest {
     stack.setStackName(stackId.getStackName());
     stack.setStackVersion(stackId.getStackVersion());
 
-    checkRequest.setRepositoryVersion(sourceRepositoryVersion);
+    checkRequest.setRepositoryVersion(sourceVersion);
     checkRequest.setSourceStackId(sourceStackId);
     checkRequest.setTargetStackId(targetStackId);
 
@@ -89,8 +92,8 @@ public class PreviousUpgradeCompletedTest {
     UpgradeEntity upgradeInProgress = Mockito.mock(UpgradeEntity.class);
     Mockito.when(upgradeInProgress.getDirection()).thenReturn(Direction.UPGRADE);
     Mockito.when(upgradeInProgress.getClusterId()).thenReturn(1L);
-    Mockito.when(upgradeInProgress.getFromVersion()).thenReturn(sourceRepositoryVersion);
-    Mockito.when(upgradeInProgress.getToVersion()).thenReturn(destRepositoryVersion);
+    Mockito.when(upgradeInProgress.getFromRepositoryVersion()).thenReturn(sourceRepositoryVersion);
+    Mockito.when(upgradeInProgress.getToRepositoryVersion()).thenReturn(targetRepositoryVersion);
 
     Mockito.when(cluster.getUpgradeInProgress()).thenReturn(upgradeInProgress);
     check = new PrerequisiteCheck(null, null);

http://git-wip-us.apache.org/repos/asf/ambari/blob/7ad307c2/ambari-server/src/test/java/org/apache/ambari/server/controller/internal/StackUpgradeConfigurationMergeTest.java
----------------------------------------------------------------------
diff --git a/ambari-server/src/test/java/org/apache/ambari/server/controller/internal/StackUpgradeConfigurationMergeTest.java b/ambari-server/src/test/java/org/apache/ambari/server/controller/internal/StackUpgradeConfigurationMergeTest.java
index 31c433d..c2d9414 100644
--- a/ambari-server/src/test/java/org/apache/ambari/server/controller/internal/StackUpgradeConfigurationMergeTest.java
+++ b/ambari-server/src/test/java/org/apache/ambari/server/controller/internal/StackUpgradeConfigurationMergeTest.java
@@ -17,6 +17,8 @@
  */
 package org.apache.ambari.server.controller.internal;
 
+import static org.easymock.EasyMock.expect;
+
 import java.util.ArrayList;
 import java.util.HashMap;
 import java.util.Map;
@@ -55,6 +57,7 @@ import org.apache.ambari.server.state.ServiceComponentHostFactory;
 import org.apache.ambari.server.state.ServiceFactory;
 import org.apache.ambari.server.state.ServiceInfo;
 import org.apache.ambari.server.state.StackId;
+import org.apache.ambari.server.state.UpgradeContext;
 import org.apache.ambari.server.state.UpgradeContextFactory;
 import org.apache.ambari.server.state.configgroup.ConfigGroupFactory;
 import org.apache.ambari.server.state.scheduler.RequestExecutionFactory;
@@ -62,6 +65,7 @@ import org.apache.ambari.server.state.stack.OsFamily;
 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.Grouping;
+import org.apache.ambari.server.state.stack.upgrade.UpgradeType;
 import org.easymock.Capture;
 import org.easymock.EasyMock;
 import org.easymock.EasyMockSupport;
@@ -81,7 +85,7 @@ import junit.framework.Assert;
 
 /**
  * Tests that
- * {@link UpgradeResourceProvider#applyStackAndProcessConfigurations(String, Cluster, String, Direction, UpgradePack, String)}
+ * {@link UpgradeResourceProvider#applyStackAndProcessConfigurations(UpgradeContext)}
  * works correctly.
  */
 public class StackUpgradeConfigurationMergeTest extends EasyMockSupport {
@@ -90,7 +94,6 @@ public class StackUpgradeConfigurationMergeTest extends EasyMockSupport {
   private static final StackId s_targetStackId = new StackId("HDP-2.5");
 
   private Injector m_injector;
-  private Clusters m_clustersMock;
   private AmbariMetaInfo m_ambariMetaInfoMock;
 
   /**
@@ -98,7 +101,6 @@ public class StackUpgradeConfigurationMergeTest extends EasyMockSupport {
    */
   @Before
   public void before() throws Exception {
-    m_clustersMock = createNiceMock(Clusters.class);
     m_ambariMetaInfoMock = createNiceMock(AmbariMetaInfo.class);
 
     MockModule mockModule = new MockModule();
@@ -154,8 +156,11 @@ public class StackUpgradeConfigurationMergeTest extends EasyMockSupport {
   @Test
   public void testMergedConfigurationsDoNotAddExplicitelyRemovedProperties() throws Exception {
     Cluster cluster = createNiceMock(Cluster.class);
-    RepositoryVersionEntity repositoryVersionEntity = createNiceMock(RepositoryVersionEntity.class);
+    RepositoryVersionEntity repositoryVersionEntity240 = createNiceMock(RepositoryVersionEntity.class);
+    RepositoryVersionEntity repositoryVersionEntity250 = createNiceMock(
+        RepositoryVersionEntity.class);
     UpgradePack upgradePack = createNiceMock(UpgradePack.class);
+    StackEntity sourceStack = createNiceMock(StackEntity.class);
     StackEntity targetStack = createNiceMock(StackEntity.class);
 
     // mocks which were bound previously
@@ -169,11 +174,16 @@ public class StackUpgradeConfigurationMergeTest extends EasyMockSupport {
     EasyMock.expect(cluster.getCurrentStackVersion()).andReturn(s_currentStackId);
     EasyMock.expect(cluster.getDesiredStackVersion()).andReturn(s_targetStackId);
 
+    EasyMock.expect(sourceStack.getStackName()).andReturn("HDP").anyTimes();
+    EasyMock.expect(sourceStack.getStackVersion()).andReturn("2.4").anyTimes();
+
     EasyMock.expect(targetStack.getStackName()).andReturn("HDP").anyTimes();
     EasyMock.expect(targetStack.getStackVersion()).andReturn("2.5").anyTimes();
 
-    EasyMock.expect(repositoryVersionEntity.getStack()).andReturn(targetStack);
-    EasyMock.expect(repositoryVersionDAO.findByStackNameAndVersion("HDP", "2.5.0.0-1234")).andReturn(repositoryVersionEntity);
+    EasyMock.expect(repositoryVersionEntity240.getStack()).andReturn(sourceStack);
+    EasyMock.expect(repositoryVersionEntity250.getStack()).andReturn(targetStack);
+    EasyMock.expect(repositoryVersionDAO.findByStackNameAndVersion("HDP", "2.4.0.0-1234")).andReturn(repositoryVersionEntity240);
+    EasyMock.expect(repositoryVersionDAO.findByStackNameAndVersion("HDP", "2.5.0.0-1234")).andReturn(repositoryVersionEntity250);
 
     EasyMock.expect(upgradePack.getGroups(Direction.UPGRADE)).andReturn(new ArrayList<Grouping>());
 
@@ -236,13 +246,22 @@ public class StackUpgradeConfigurationMergeTest extends EasyMockSupport {
 
     EasyMock.expectLastCall();
 
+    UpgradeContext context = createNiceMock(UpgradeContext.class);
+    expect(context.getCluster()).andReturn(cluster).atLeastOnce();
+    expect(context.getType()).andReturn(UpgradeType.ROLLING).atLeastOnce();
+    expect(context.getDirection()).andReturn(Direction.UPGRADE).atLeastOnce();
+    expect(context.getSourceRepositoryVersion()).andReturn(repositoryVersionEntity240).atLeastOnce();
+    expect(context.getTargetRepositoryVersion()).andReturn(repositoryVersionEntity250).atLeastOnce();
+    expect(context.getUpgradePack()).andReturn(upgradePack).atLeastOnce();
+    expect(context.getHostRoleCommandFactory()).andStubReturn(m_injector.getInstance(HostRoleCommandFactory.class));
+    expect(context.getRoleGraphFactory()).andStubReturn(m_injector.getInstance(RoleGraphFactory.class));
+
     replayAll();
 
     UpgradeResourceProvider upgradeResourceProvider = new UpgradeResourceProvider(amc);
     m_injector.injectMembers(upgradeResourceProvider);
 
-    upgradeResourceProvider.applyStackAndProcessConfigurations("HDP", cluster, "2.5.0.0-1234",
-        Direction.UPGRADE, upgradePack, "admin");
+    upgradeResourceProvider.applyStackAndProcessConfigurations(context);
 
     // assertion time!
     Map<String, Map<String, String>> mergedConfigurations = capturedArgument.getValue();

http://git-wip-us.apache.org/repos/asf/ambari/blob/7ad307c2/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 1e888c1..9652c69 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
@@ -95,6 +95,7 @@ import org.apache.ambari.server.state.Service;
 import org.apache.ambari.server.state.ServiceComponent;
 import org.apache.ambari.server.state.ServiceComponentHost;
 import org.apache.ambari.server.state.StackId;
+import org.apache.ambari.server.state.UpgradeContext;
 import org.apache.ambari.server.state.UpgradeHelper;
 import org.apache.ambari.server.state.UpgradeState;
 import org.apache.ambari.server.state.stack.UpgradePack;
@@ -141,6 +142,10 @@ public class UpgradeResourceProviderTest {
   private ConfigFactory configFactory;
   private HostRoleCommandDAO hrcDAO;
 
+  RepositoryVersionEntity repoVersionEntity2110;
+  RepositoryVersionEntity repoVersionEntity2111;
+  RepositoryVersionEntity repoVersionEntity2200;
+
   @Before
   public void before() throws Exception {
     SecurityContextHolder.getContext().setAuthentication(
@@ -199,26 +204,27 @@ public class UpgradeResourceProviderTest {
     StackId stack211 = new StackId("HDP-2.1.1");
     StackId stack220 = new StackId("HDP-2.2.0");
 
-    RepositoryVersionEntity repoVersionEntity = new RepositoryVersionEntity();
-    repoVersionEntity.setDisplayName("My New Version 1");
-    repoVersionEntity.setOperatingSystems("");
-    repoVersionEntity.setStack(stackEntity211);
-    repoVersionEntity.setVersion("2.1.1.0");
-    repoVersionDao.create(repoVersionEntity);
 
-    repoVersionEntity = new RepositoryVersionEntity();
-    repoVersionEntity.setDisplayName("My New Version 2 for patch upgrade");
-    repoVersionEntity.setOperatingSystems("");
-    repoVersionEntity.setStack(stackEntity211);
-    repoVersionEntity.setVersion("2.1.1.1");
-    repoVersionDao.create(repoVersionEntity);
+    repoVersionEntity2110 = new RepositoryVersionEntity();
+    repoVersionEntity2110.setDisplayName("My New Version 1");
+    repoVersionEntity2110.setOperatingSystems("");
+    repoVersionEntity2110.setStack(stackEntity211);
+    repoVersionEntity2110.setVersion("2.1.1.0");
+    repoVersionDao.create(repoVersionEntity2110);
 
-    repoVersionEntity = new RepositoryVersionEntity();
-    repoVersionEntity.setDisplayName("My New Version 3 for major upgrade");
-    repoVersionEntity.setOperatingSystems("");
-    repoVersionEntity.setStack(stackEntity220);
-    repoVersionEntity.setVersion("2.2.0.0");
-    repoVersionDao.create(repoVersionEntity);
+    repoVersionEntity2111 = new RepositoryVersionEntity();
+    repoVersionEntity2111.setDisplayName("My New Version 2 for patch upgrade");
+    repoVersionEntity2111.setOperatingSystems("");
+    repoVersionEntity2111.setStack(stackEntity211);
+    repoVersionEntity2111.setVersion("2.1.1.1");
+    repoVersionDao.create(repoVersionEntity2111);
+
+    repoVersionEntity2200 = new RepositoryVersionEntity();
+    repoVersionEntity2200.setDisplayName("My New Version 3 for major upgrade");
+    repoVersionEntity2200.setOperatingSystems("");
+    repoVersionEntity2200.setStack(stackEntity220);
+    repoVersionEntity2200.setVersion("2.2.0.0");
+    repoVersionDao.create(repoVersionEntity2200);
 
     clusters = injector.getInstance(Clusters.class);
 
@@ -603,7 +609,7 @@ public class UpgradeResourceProviderTest {
     // this should get skipped
     ServiceComponent component = service.getServiceComponent("ZOOKEEPER_SERVER");
     ServiceComponentHost sch = component.addServiceComponentHost("h2");
-    sch.setVersion("2.2.2.2");
+    sch.setVersion(repoVersionEntity2200.getVersion());
 
     // start out with 0 (sanity check)
     List<UpgradeEntity> upgrades = upgradeDao.findUpgrades(cluster.getClusterId());
@@ -621,8 +627,8 @@ public class UpgradeResourceProviderTest {
     UpgradeEntity upgradeEntity = new UpgradeEntity();
     upgradeEntity.setClusterId(cluster.getClusterId());
     upgradeEntity.setDirection(Direction.UPGRADE);
-    upgradeEntity.setFromVersion("2.1.1.1");
-    upgradeEntity.setToVersion("2.2.2.2");
+    upgradeEntity.setFromRepositoryVersion(repoVersionEntity2111);
+    upgradeEntity.setToRepositoryVersion(repoVersionEntity2200);
     upgradeEntity.setUpgradePackage("upgrade_test");
     upgradeEntity.setUpgradeType(UpgradeType.ROLLING);
     upgradeEntity.setRequestEntity(requestEntity);
@@ -636,7 +642,7 @@ public class UpgradeResourceProviderTest {
 
     Map<String, Object> requestProps = new HashMap<>();
     requestProps.put(UpgradeResourceProvider.UPGRADE_CLUSTER_NAME, "c1");
-    requestProps.put(UpgradeResourceProvider.UPGRADE_REPO_VERSION, "2.1.1.1");
+    requestProps.put(UpgradeResourceProvider.UPGRADE_REPO_ID, 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());
@@ -720,8 +726,8 @@ public class UpgradeResourceProviderTest {
 
     UpgradeEntity entity = upgradeDao.findUpgrade(Long.parseLong(id));
     assertNotNull(entity);
-    assertEquals("2.1.1", entity.getFromVersion());
-    assertEquals("2.2.0.0", entity.getToVersion());
+    assertEquals("2.1.1.0", entity.getFromRepositoryVersion().getVersion());
+    assertEquals("2.2.0.0", entity.getToRepositoryVersion().getVersion());
     assertEquals(Direction.DOWNGRADE, entity.getDirection());
 
     StageDAO dao = injector.getInstance(StageDAO.class);
@@ -1186,7 +1192,7 @@ public class UpgradeResourceProviderTest {
         configHelper.getDefaultProperties(EasyMock.eq(stack220), EasyMock.anyObject(Cluster.class))).andReturn(
         stack220Configs).anyTimes();
 
-    Capture<Map<String, Map<String, String>>> expectedConfigurationsCapture = new Capture<Map<String, Map<String, String>>>();
+    Capture<Map<String, Map<String, String>>> expectedConfigurationsCapture = new Capture<>();
 
     configHelper.createConfigTypes(EasyMock.anyObject(Cluster.class),
         EasyMock.anyObject(AmbariManagementController.class),
@@ -1200,8 +1206,17 @@ public class UpgradeResourceProviderTest {
     UpgradeResourceProvider upgradeResourceProvider = createProvider(amc);
 
     Map<String, UpgradePack> upgradePacks = ambariMetaInfo.getUpgradePacks("HDP", "2.1.1");
-    UpgradePack upgrade = upgradePacks.get("upgrade_to_new_stack");
-    upgradeResourceProvider.applyStackAndProcessConfigurations(stack211.getStackName(), cluster, "2.2.0.0", Direction.UPGRADE, upgrade, "admin");
+    UpgradePack upgradePack = upgradePacks.get("upgrade_to_new_stack");
+
+    UpgradeContext context = EasyMock.createNiceMock(UpgradeContext.class);
+    expect(context.getCluster()).andReturn(cluster).atLeastOnce();
+    expect(context.getType()).andReturn(UpgradeType.ROLLING).atLeastOnce();
+    expect(context.getDirection()).andReturn(Direction.UPGRADE).atLeastOnce();
+    expect(context.getSourceRepositoryVersion()).andReturn(repoVersionEntity2111).atLeastOnce();
+    expect(context.getTargetRepositoryVersion()).andReturn(repoVersionEntity2200).atLeastOnce();
+    expect(context.getUpgradePack()).andReturn(upgradePack).atLeastOnce();
+
+    upgradeResourceProvider.applyStackAndProcessConfigurations(context);
 
     Map<String, Map<String, String>> expectedConfigurations = expectedConfigurationsCapture.getValue();
     Map<String, String> expectedFooType = expectedConfigurations.get("foo-site");
@@ -1268,7 +1283,7 @@ public class UpgradeResourceProviderTest {
       Map<String, String> map = gson.<Map<String, String>> fromJson(se.getCommandParamsStage(),Map.class);
       assertTrue(map.containsKey("upgrade_direction"));
       assertEquals("upgrade", map.get("upgrade_direction"));
-      
+
       if(map.containsKey("upgrade_type")){
         assertEquals("rolling_upgrade", map.get("upgrade_type"));
       }

http://git-wip-us.apache.org/repos/asf/ambari/blob/7ad307c2/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 f1416d2..f63825b 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
@@ -169,7 +169,7 @@ public class UpgradeSummaryResourceProviderTest {
 
     clusters.addHost("h1");
     Host host = clusters.getHost("h1");
-    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.4");
     host.setHostAttributes(hostAttributes);
@@ -280,8 +280,16 @@ public class UpgradeSummaryResourceProviderTest {
     upgrade.setUpgradePackage("some-name");
     upgrade.setUpgradeType(UpgradeType.ROLLING);
     upgrade.setDirection(Direction.UPGRADE);
-    upgrade.setFromVersion("2.2.0.0");
-    upgrade.setToVersion("2.2.0.1");
+
+    RepositoryVersionEntity repositoryVersion2200 = injector.getInstance(
+        RepositoryVersionDAO.class).findByStackNameAndVersion("HDP", "2.2.0.0-1234");
+
+    RepositoryVersionEntity repositoryVersion2201 = injector.getInstance(
+        RepositoryVersionDAO.class).findByStackNameAndVersion("HDP", "2.2.0.1-1234");
+
+    upgrade.setFromRepositoryVersion(repositoryVersion2200);
+    upgrade.setToRepositoryVersion(repositoryVersion2201);
+
     upgradeDAO.create(upgrade);
 
     // Resource used to make assertions.

http://git-wip-us.apache.org/repos/asf/ambari/blob/7ad307c2/ambari-server/src/test/java/org/apache/ambari/server/orm/dao/UpgradeDAOTest.java
----------------------------------------------------------------------
diff --git a/ambari-server/src/test/java/org/apache/ambari/server/orm/dao/UpgradeDAOTest.java b/ambari-server/src/test/java/org/apache/ambari/server/orm/dao/UpgradeDAOTest.java
index 78fedfe..d712b03 100644
--- a/ambari-server/src/test/java/org/apache/ambari/server/orm/dao/UpgradeDAOTest.java
+++ b/ambari-server/src/test/java/org/apache/ambari/server/orm/dao/UpgradeDAOTest.java
@@ -33,11 +33,13 @@ import org.apache.ambari.server.actionmanager.HostRoleStatus;
 import org.apache.ambari.server.orm.GuiceJpaInitializer;
 import org.apache.ambari.server.orm.InMemoryDefaultTestModule;
 import org.apache.ambari.server.orm.OrmTestHelper;
+import org.apache.ambari.server.orm.entities.RepositoryVersionEntity;
 import org.apache.ambari.server.orm.entities.RequestEntity;
 import org.apache.ambari.server.orm.entities.StageEntity;
 import org.apache.ambari.server.orm.entities.UpgradeEntity;
 import org.apache.ambari.server.orm.entities.UpgradeGroupEntity;
 import org.apache.ambari.server.orm.entities.UpgradeItemEntity;
+import org.apache.ambari.server.state.StackId;
 import org.apache.ambari.server.state.UpgradeState;
 import org.apache.ambari.server.state.stack.upgrade.Direction;
 import org.apache.ambari.server.state.stack.upgrade.UpgradeType;
@@ -62,6 +64,10 @@ public class UpgradeDAOTest {
 
   private OrmTestHelper helper;
 
+  RepositoryVersionEntity repositoryVersion2200;
+  RepositoryVersionEntity repositoryVersion2500;
+  RepositoryVersionEntity repositoryVersion2511;
+
   /**
    *
    */
@@ -83,12 +89,16 @@ public class UpgradeDAOTest {
     requestEntity.setStages(new ArrayList<StageEntity>());
     requestDAO.create(requestEntity);
 
+    repositoryVersion2200 = helper.getOrCreateRepositoryVersion(new StackId("HDP", "2.2.0"), "2.2.0.0-1234");
+    repositoryVersion2500 = helper.getOrCreateRepositoryVersion(new StackId("HDP", "2.5.0"), "2.5.0.0-4567");
+    repositoryVersion2511 = helper.getOrCreateRepositoryVersion(new StackId("HDP", "2.5.0"), "2.5.1.1-4567");
+
     // create upgrade entities
     UpgradeEntity entity = new UpgradeEntity();
     entity.setClusterId(clusterId.longValue());
     entity.setRequestEntity(requestEntity);
-    entity.setFromVersion("");
-    entity.setToVersion("");
+    entity.setToRepositoryVersion(repositoryVersion2200);
+    entity.setFromRepositoryVersion(repositoryVersion2500);
     entity.setUpgradeType(UpgradeType.ROLLING);
     entity.setUpgradePackage("test-upgrade");
     entity.setDowngradeAllowed(true);
@@ -98,7 +108,7 @@ public class UpgradeDAOTest {
     group.setTitle("group title");
 
     // create 2 items
-    List<UpgradeItemEntity> items = new ArrayList<UpgradeItemEntity>();
+    List<UpgradeItemEntity> items = new ArrayList<>();
     UpgradeItemEntity item = new UpgradeItemEntity();
     item.setState(UpgradeState.IN_PROGRESS);
     item.setStageId(Long.valueOf(1L));
@@ -165,8 +175,8 @@ public class UpgradeDAOTest {
     entity1.setClusterId(clusterId.longValue());
     entity1.setDirection(Direction.UPGRADE);
     entity1.setRequestEntity(requestEntity);
-    entity1.setFromVersion("2.2.0.0-1234");
-    entity1.setToVersion("2.3.0.0-4567");
+    entity1.setToRepositoryVersion(repositoryVersion2200);
+    entity1.setFromRepositoryVersion(repositoryVersion2500);
     entity1.setUpgradeType(UpgradeType.ROLLING);
     entity1.setUpgradePackage("test-upgrade");
     entity1.setDowngradeAllowed(true);
@@ -176,8 +186,8 @@ public class UpgradeDAOTest {
     entity2.setClusterId(clusterId.longValue());
     entity2.setDirection(Direction.DOWNGRADE);
     entity2.setRequestEntity(requestEntity);
-    entity2.setFromVersion("2.3.0.0-4567");
-    entity2.setToVersion("2.2.0.0-1234");
+    entity2.setToRepositoryVersion(repositoryVersion2500);
+    entity2.setFromRepositoryVersion(repositoryVersion2200);
     entity2.setUpgradeType(UpgradeType.ROLLING);
     entity2.setUpgradePackage("test-upgrade");
     entity2.setDowngradeAllowed(true);
@@ -187,8 +197,8 @@ public class UpgradeDAOTest {
     entity3.setClusterId(clusterId.longValue());
     entity3.setDirection(Direction.UPGRADE);
     entity3.setRequestEntity(requestEntity);
-    entity3.setFromVersion("2.2.0.0-1234");
-    entity3.setToVersion("2.3.1.1-4567");
+    entity3.setToRepositoryVersion(repositoryVersion2200);
+    entity3.setFromRepositoryVersion(repositoryVersion2500);
     entity3.setUpgradeType(UpgradeType.ROLLING);
     entity3.setUpgradePackage("test-upgrade");
     entity3.setDowngradeAllowed(true);
@@ -217,8 +227,8 @@ public class UpgradeDAOTest {
     upgradeEntity.setClusterId(clusterId.longValue());
     upgradeEntity.setDirection(Direction.UPGRADE);
     upgradeEntity.setRequestEntity(requestEntity);
-    upgradeEntity.setFromVersion("2.2.0.0-1234");
-    upgradeEntity.setToVersion("2.3.0.0-4567");
+    upgradeEntity.setToRepositoryVersion(repositoryVersion2200);
+    upgradeEntity.setFromRepositoryVersion(repositoryVersion2500);
     upgradeEntity.setUpgradeType(UpgradeType.ROLLING);
     upgradeEntity.setUpgradePackage("test-upgrade");
     dao.create(upgradeEntity);

http://git-wip-us.apache.org/repos/asf/ambari/blob/7ad307c2/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 b95a447..d875443 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
@@ -50,8 +50,6 @@ import org.apache.ambari.server.orm.GuiceJpaInitializer;
 import org.apache.ambari.server.orm.InMemoryDefaultTestModule;
 import org.apache.ambari.server.orm.OrmTestHelper;
 import org.apache.ambari.server.orm.dao.ClusterVersionDAO;
-import org.apache.ambari.server.orm.dao.HostComponentDesiredStateDAO;
-import org.apache.ambari.server.orm.dao.HostComponentStateDAO;
 import org.apache.ambari.server.orm.dao.HostDAO;
 import org.apache.ambari.server.orm.dao.HostVersionDAO;
 import org.apache.ambari.server.orm.dao.RepositoryVersionDAO;
@@ -83,7 +81,6 @@ 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;
@@ -147,10 +144,6 @@ public class UpgradeActionTest {
   @Inject
   private ServiceComponentDesiredStateDAO serviceComponentDesiredStateDAO;
   @Inject
-  private HostComponentDesiredStateDAO hostComponentDesiredStateDAO;
-  @Inject
-  private HostComponentStateDAO hostComponentStateDAO;
-  @Inject
   private StackDAO stackDAO;
   @Inject
   private AmbariMetaInfo ambariMetaInfo;
@@ -159,6 +152,10 @@ public class UpgradeActionTest {
   @Inject
   private ConfigFactory configFactory;
 
+  private RepositoryVersionEntity repositoryVersion2110;
+  private RepositoryVersionEntity repositoryVersion2111;
+  private RepositoryVersionEntity repositoryVersion2201;
+
   @Before
   public void setup() throws Exception {
     m_injector = Guice.createInjector(new InMemoryDefaultTestModule());
@@ -173,6 +170,10 @@ public class UpgradeActionTest {
     Field field = AmbariServer.class.getDeclaredField("clusterController");
     field.setAccessible(true);
     field.set(null, amc);
+
+    repositoryVersion2110 = m_helper.getOrCreateRepositoryVersion(HDP_21_STACK, HDP_2_1_1_0);
+    repositoryVersion2111 = m_helper.getOrCreateRepositoryVersion(HDP_21_STACK, HDP_2_1_1_1);
+    repositoryVersion2201 = m_helper.getOrCreateRepositoryVersion(HDP_22_STACK, HDP_2_2_0_1);
   }
 
   @After
@@ -193,7 +194,7 @@ public class UpgradeActionTest {
 
     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);
@@ -230,7 +231,7 @@ public class UpgradeActionTest {
 
     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);
@@ -293,7 +294,7 @@ public class UpgradeActionTest {
 
     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);
@@ -379,7 +380,7 @@ public class UpgradeActionTest {
     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);
@@ -468,10 +469,7 @@ public class UpgradeActionTest {
     Collection<Config> configs = cluster.getAllConfigs();
     Assert.assertFalse(configs.isEmpty());
 
-    Map<String, String> commandParams = new HashMap<String, String>();
-    commandParams.put(UpdateDesiredStackAction.COMMAND_PARAM_ORIGINAL_STACK, sourceStack.getStackId());
-    commandParams.put(UpdateDesiredStackAction.COMMAND_PARAM_TARGET_STACK, targetStack.getStackId());
-    commandParams.put(UpdateDesiredStackAction.COMMAND_PARAM_DIRECTION, Direction.UPGRADE.toString());
+    Map<String, String> commandParams = new HashMap<>();
     commandParams.put(UpdateDesiredStackAction.COMMAND_PARAM_VERSION, targetRepo);
     commandParams.put(UpdateDesiredStackAction.COMMAND_PARAM_UPGRADE_PACK, upgradePackName);
 
@@ -514,7 +512,7 @@ public class UpgradeActionTest {
 
     makeDowngradeCluster(sourceStack, sourceRepo, targetStack, targetRepo);
 
-    Map<String, String> commandParams = new HashMap<String, String>();
+    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());
@@ -572,7 +570,7 @@ public class UpgradeActionTest {
 
     makeTwoUpgradesWhereLastDidNotComplete(sourceStack, sourceRepo, midStack, midRepo, targetStack, targetRepo);
 
-    Map<String, String> commandParams = new HashMap<String, String>();
+    Map<String, String> commandParams = new HashMap<>();
     commandParams.put(FinalizeUpgradeAction.UPGRADE_DIRECTION_KEY, "downgrade");
     commandParams.put(FinalizeUpgradeAction.VERSION_KEY, midRepo);
     commandParams.put(FinalizeUpgradeAction.ORIGINAL_STACK_KEY, sourceStack.getStackId());
@@ -615,7 +613,7 @@ public class UpgradeActionTest {
     verifyBaseRepoURL(helper, cluster, host, HDP_211_CENTOS6_REPO_URL);
 
     // 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);
 
@@ -679,7 +677,7 @@ public class UpgradeActionTest {
     verifyBaseRepoURL(helper, cluster, host, HDP_211_CENTOS6_REPO_URL);
 
     // 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);
 
@@ -728,7 +726,7 @@ public class UpgradeActionTest {
     cluster.setCurrentStackVersion(sourceStack);
     cluster.setDesiredStackVersion(targetStack);
 
-    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());
@@ -801,7 +799,7 @@ public class UpgradeActionTest {
     Collection<Config> configs = cluster.getAllConfigs();
     assertEquals(8, configs.size());
 
-    Map<String, String> commandParams = new HashMap<String, String>();
+    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());
@@ -908,7 +906,7 @@ public class UpgradeActionTest {
 
     // 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());
@@ -942,8 +940,8 @@ public class UpgradeActionTest {
   public void testUpgradeHistory() throws Exception {
     StackId sourceStack = HDP_21_STACK;
     StackId targetStack = HDP_21_STACK;
-    String sourceRepo = HDP_2_1_1_0;
-    String targetRepo = HDP_2_1_1_1;
+    String sourceRepo = repositoryVersion2110.getVersion();
+    String targetRepo = repositoryVersion2111.getVersion();
     String hostName = "h1";
 
     createUpgradeClusterAndSourceRepo(sourceStack, sourceRepo, hostName);
@@ -982,8 +980,8 @@ public class UpgradeActionTest {
     upgradeEntity.setClusterId(cluster.getClusterId());
     upgradeEntity.setRequestEntity(requestEntity);
     upgradeEntity.setUpgradePackage("");
-    upgradeEntity.setFromVersion(sourceRepo);
-    upgradeEntity.setToVersion(targetRepo);
+    upgradeEntity.setFromRepositoryVersion(repositoryVersion2110);
+    upgradeEntity.setToRepositoryVersion(repositoryVersion2111);
     upgradeEntity.setUpgradeType(UpgradeType.NON_ROLLING);
     upgradeDAO.create(upgradeEntity);
 
@@ -1000,7 +998,7 @@ public class UpgradeActionTest {
 
     // Finalize the upgrade, passing in the request ID so that history is
     // created
-    Map<String, String> commandParams = new HashMap<String, String>();
+    Map<String, String> commandParams = new HashMap<>();
     commandParams.put(FinalizeUpgradeAction.REQUEST_ID, String.valueOf(requestEntity.getRequestId()));
     commandParams.put(FinalizeUpgradeAction.UPGRADE_DIRECTION_KEY, "upgrade");
     commandParams.put(FinalizeUpgradeAction.VERSION_KEY, targetRepo);
@@ -1074,8 +1072,8 @@ public class UpgradeActionTest {
   }
 
   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/7ad307c2/ambari-server/src/test/java/org/apache/ambari/server/state/ServiceComponentTest.java
----------------------------------------------------------------------
diff --git a/ambari-server/src/test/java/org/apache/ambari/server/state/ServiceComponentTest.java b/ambari-server/src/test/java/org/apache/ambari/server/state/ServiceComponentTest.java
index 135edb6..780398c 100644
--- a/ambari-server/src/test/java/org/apache/ambari/server/state/ServiceComponentTest.java
+++ b/ambari-server/src/test/java/org/apache/ambari/server/state/ServiceComponentTest.java
@@ -187,7 +187,7 @@ public class ServiceComponentTest {
     h.setIPv4(hostname + "ipv4");
     h.setIPv6(hostname + "ipv6");
 
-    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.3");
     h.setHostAttributes(hostAttributes);
@@ -286,7 +286,7 @@ public class ServiceComponentTest {
     sch.setState(State.INSTALLED);
 
     Map<String, ServiceComponentHost> compHosts =
-      new HashMap<String, ServiceComponentHost>();
+      new HashMap<>();
     compHosts.put("h1", sch);
     component.addServiceComponentHosts(compHosts);
     Assert.assertEquals(1, component.getServiceComponentHosts().size());
@@ -361,9 +361,9 @@ public class ServiceComponentTest {
     sc.setDesiredState(State.INSTALLED);
     Assert.assertEquals(State.INSTALLED, sc.getDesiredState());
 
-    sc.setDesiredStackVersion(new StackId("HDP-2.2.0"));
-    StackId stackId = sc.getDesiredStackVersion();
-    Assert.assertEquals(new StackId("HDP", "2.2.0"), stackId);
+    StackId stackId = new StackId("HDP-2.2.0");
+    sc.setDesiredStackVersion(stackId);
+    Assert.assertEquals(stackId, sc.getDesiredStackVersion());
 
     Assert.assertEquals("HDP-2.2.0", sc.getDesiredStackVersion().getStackId());
 
@@ -372,7 +372,7 @@ public class ServiceComponentTest {
 
     Assert.assertNotNull(serviceComponentDesiredStateEntity);
 
-    UpgradeEntity upgradeEntity = createUpgradeEntity("2.2.0.0", "2.2.0.1");
+    UpgradeEntity upgradeEntity = createUpgradeEntity(stackId, "2.2.0.0", "2.2.0.1");
     ServiceComponentHistoryEntity history = new ServiceComponentHistoryEntity();
     history.setFromStack(serviceComponentDesiredStateEntity.getDesiredStack());
     history.setToStack(serviceComponentDesiredStateEntity.getDesiredStack());
@@ -430,7 +430,7 @@ public class ServiceComponentTest {
         serviceComponentHostFactory.createNew(sc, "h2");
 
     Map<String, ServiceComponentHost> compHosts =
-        new HashMap<String, ServiceComponentHost>();
+        new HashMap<>();
     compHosts.put("h1", sch1);
     compHosts.put("h2", sch2);
     sc.addServiceComponentHosts(compHosts);
@@ -478,9 +478,9 @@ public class ServiceComponentTest {
     sc.setDesiredState(State.INSTALLED);
     Assert.assertEquals(State.INSTALLED, sc.getDesiredState());
 
-    sc.setDesiredStackVersion(new StackId("HDP-2.2.0"));
-    StackId stackId = sc.getDesiredStackVersion();
-    Assert.assertEquals(new StackId("HDP", "2.2.0"), stackId);
+    StackId stackId = new StackId("HDP-2.2.0");
+    sc.setDesiredStackVersion(stackId);
+    Assert.assertEquals(stackId, sc.getDesiredStackVersion());
 
     Assert.assertEquals("HDP-2.2.0", sc.getDesiredStackVersion().getStackId());
 
@@ -490,7 +490,7 @@ public class ServiceComponentTest {
 
     Assert.assertNotNull(serviceComponentDesiredStateEntity);
 
-    UpgradeEntity upgradeEntity = createUpgradeEntity("2.2.0.0", "2.2.0.1");
+    UpgradeEntity upgradeEntity = createUpgradeEntity(stackId, "2.2.0.0", "2.2.0.1");
     ServiceComponentHistoryEntity history = new ServiceComponentHistoryEntity();
     history.setFromStack(serviceComponentDesiredStateEntity.getDesiredStack());
     history.setToStack(serviceComponentDesiredStateEntity.getDesiredStack());
@@ -647,7 +647,10 @@ public class ServiceComponentTest {
    * @param toVersion
    * @return
    */
-  private UpgradeEntity createUpgradeEntity(String fromVersion, String toVersion) {
+  private UpgradeEntity createUpgradeEntity(StackId stackId, String fromVersion, String toVersion) {
+    RepositoryVersionEntity toRepositoryVersion = helper.getOrCreateRepositoryVersion(stackId, fromVersion);
+    RepositoryVersionEntity fromRepositoryVersion = helper.getOrCreateRepositoryVersion(stackId, fromVersion);
+
     RequestDAO requestDAO = injector.getInstance(RequestDAO.class);
     RequestEntity requestEntity = new RequestEntity();
     requestEntity.setRequestId(99L);
@@ -660,8 +663,8 @@ public class ServiceComponentTest {
     UpgradeEntity upgradeEntity = new UpgradeEntity();
     upgradeEntity.setClusterId(cluster.getClusterId());
     upgradeEntity.setDirection(Direction.UPGRADE);
-    upgradeEntity.setFromVersion(fromVersion);
-    upgradeEntity.setToVersion(toVersion);
+    upgradeEntity.setFromRepositoryVersion(fromRepositoryVersion);
+    upgradeEntity.setToRepositoryVersion(toRepositoryVersion);
     upgradeEntity.setUpgradePackage("upgrade_test");
     upgradeEntity.setUpgradeType(UpgradeType.ROLLING);
     upgradeEntity.setRequestEntity(requestEntity);

http://git-wip-us.apache.org/repos/asf/ambari/blob/7ad307c2/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 0dd7f58..b5dec7f 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
@@ -47,6 +47,7 @@ import org.apache.ambari.server.controller.ConfigurationRequest;
 import org.apache.ambari.server.orm.GuiceJpaInitializer;
 import org.apache.ambari.server.orm.InMemoryDefaultTestModule;
 import org.apache.ambari.server.orm.OrmTestHelper;
+import org.apache.ambari.server.orm.entities.RepositoryVersionEntity;
 import org.apache.ambari.server.security.TestAuthenticationFactory;
 import org.apache.ambari.server.security.authorization.AuthorizationException;
 import org.apache.ambari.server.stack.HostsType;
@@ -95,8 +96,8 @@ import com.google.inject.util.Modules;
  */
 public class UpgradeHelperTest {
 
-  private static final StackId HDP_21 = new StackId("HDP-2.1.1");
-  private static final StackId HDP_22 = new StackId("HDP-2.2.0");
+  private static final StackId STACK_ID_HDP_211 = new StackId("HDP-2.1.1");
+  private static final StackId STACK_ID_HDP_220 = new StackId("HDP-2.2.0");
   private static final String UPGRADE_VERSION = "2.2.1.0-1234";
   private static final String DOWNGRADE_VERSION = "2.2.0.0-1234";
 
@@ -109,7 +110,10 @@ public class UpgradeHelperTest {
   private ConfigHelper m_configHelper;
   private AmbariManagementController m_managementController;
   private Gson m_gson = new Gson();
-  private UpgradeContextFactory m_upgradeContextFactory;
+
+  private RepositoryVersionEntity repositoryVersion2110;
+  private RepositoryVersionEntity repositoryVersion2200;
+  private RepositoryVersionEntity repositoryVersion2210;
 
   /**
    * Because test cases need to share config mocks, put common ones in this function.
@@ -153,7 +157,10 @@ public class UpgradeHelperTest {
     m_upgradeHelper = injector.getInstance(UpgradeHelper.class);
     m_masterHostResolver = EasyMock.createMock(MasterHostResolver.class);
     m_managementController = injector.getInstance(AmbariManagementController.class);
-    m_upgradeContextFactory = injector.getInstance(UpgradeContextFactory.class);
+
+    repositoryVersion2110 = helper.getOrCreateRepositoryVersion(STACK_ID_HDP_211, "2.1.1.0-1234");
+    repositoryVersion2200 = helper.getOrCreateRepositoryVersion(STACK_ID_HDP_220, DOWNGRADE_VERSION);
+    repositoryVersion2210 = helper.getOrCreateRepositoryVersion(STACK_ID_HDP_220, UPGRADE_VERSION);
 
     // Set the authenticated user
     // TODO: remove this or replace the authenticated user to test authorization rules
@@ -204,10 +211,8 @@ public class UpgradeHelperTest {
 
     Cluster cluster = makeCluster();
 
-    UpgradeContext context = m_upgradeContextFactory.create(cluster, UpgradeType.ROLLING,
-        Direction.UPGRADE, UPGRADE_VERSION, new HashMap<String, Object>());
-
-    context.setResolver(m_masterHostResolver);
+    UpgradeContext context = getMockUpgradeContext(cluster, Direction.UPGRADE, UpgradeType.ROLLING,
+        repositoryVersion2210);
 
     List<UpgradeGroupHolder> groups = m_upgradeHelper.createSequence(upgrade, context);
 
@@ -285,12 +290,14 @@ public class UpgradeHelperTest {
 
     Cluster cluster = makeCluster();
 
-    UpgradeContext context = m_upgradeContextFactory.create(cluster, UpgradeType.ROLLING,
-        Direction.UPGRADE, UPGRADE_VERSION, new HashMap<String, Object>());
-    context.setResolver(m_masterHostResolver);
+    UpgradeContext context = getMockUpgradeContext(cluster, Direction.UPGRADE, UpgradeType.ROLLING,
+        repositoryVersion2210, m_masterHostResolver, false);
+
     context.setSupportedServices(Collections.singleton("ZOOKEEPER"));
     context.setScope(UpgradeScope.PARTIAL);
 
+    replay(context);
+
     List<Grouping> groupings = upgrade.getGroups(Direction.UPGRADE);
     assertEquals(8, groupings.size());
     assertEquals(UpgradeScope.COMPLETE, groupings.get(6).scope);
@@ -343,13 +350,14 @@ public class UpgradeHelperTest {
 
     Cluster cluster = makeCluster();
 
-    UpgradeContext context = m_upgradeContextFactory.create(cluster, UpgradeType.ROLLING,
-        Direction.UPGRADE, UPGRADE_VERSION, new HashMap<String, Object>());
+    UpgradeContext context = getMockUpgradeContext(cluster, Direction.UPGRADE, UpgradeType.ROLLING,
+        repositoryVersion2210, m_masterHostResolver, false);
 
-    context.setResolver(m_masterHostResolver);
     context.setSupportedServices(Collections.singleton("ZOOKEEPER"));
     context.setScope(UpgradeScope.COMPLETE);
 
+    replay(context);
+
     List<Grouping> groupings = upgrade.getGroups(Direction.UPGRADE);
     assertEquals(8, groupings.size());
     assertEquals(UpgradeScope.COMPLETE, groupings.get(6).scope);
@@ -400,9 +408,8 @@ public class UpgradeHelperTest {
 
     Cluster cluster = makeCluster();
 
-    UpgradeContext context = m_upgradeContextFactory.create(cluster, UpgradeType.ROLLING,
-        Direction.UPGRADE, UPGRADE_VERSION, new HashMap<String, Object>());
-    context.setResolver(m_masterHostResolver);
+    UpgradeContext context = getMockUpgradeContext(cluster, Direction.UPGRADE, UpgradeType.ROLLING,
+        repositoryVersion2210, m_masterHostResolver, true);
 
     List<UpgradeGroupHolder> groups = m_upgradeHelper.createSequence(upgrade, context);
 
@@ -453,9 +460,8 @@ public class UpgradeHelperTest {
     // use a "real" master host resolver here so that we can actually test MM
     MasterHostResolver masterHostResolver = new MasterHostResolver(null, cluster, "");
 
-    UpgradeContext context = m_upgradeContextFactory.create(cluster, UpgradeType.ROLLING,
-        Direction.UPGRADE, UPGRADE_VERSION, new HashMap<String, Object>());
-    context.setResolver(masterHostResolver);
+    UpgradeContext context = getMockUpgradeContext(cluster, Direction.UPGRADE, UpgradeType.ROLLING,
+        repositoryVersion2210, masterHostResolver, true);
 
     List<UpgradeGroupHolder> groups = m_upgradeHelper.createSequence(upgrade, context);
     assertEquals(7, groups.size());
@@ -485,9 +491,8 @@ public class UpgradeHelperTest {
 
     Cluster cluster = makeCluster();
 
-    UpgradeContext context = m_upgradeContextFactory.create(cluster, UpgradeType.ROLLING,
-        Direction.UPGRADE, UPGRADE_VERSION, new HashMap<String, Object>());
-    context.setResolver(m_masterHostResolver);
+    UpgradeContext context = getMockUpgradeContext(cluster, Direction.UPGRADE, UpgradeType.ROLLING,
+        repositoryVersion2210);
 
     List<UpgradeGroupHolder> groups = m_upgradeHelper.createSequence(upgrade, context);
 
@@ -535,9 +540,8 @@ public class UpgradeHelperTest {
     assertEquals(1, schs.size());
     assertEquals(HostState.HEARTBEAT_LOST, schs.get(0).getHostState());
 
-    UpgradeContext context = m_upgradeContextFactory.create(cluster, UpgradeType.ROLLING,
-        Direction.UPGRADE, UPGRADE_VERSION, new HashMap<String, Object>());
-    context.setResolver(m_masterHostResolver);
+    UpgradeContext context = getMockUpgradeContext(cluster, Direction.UPGRADE, UpgradeType.ROLLING,
+        repositoryVersion2210);
 
     List<UpgradeGroupHolder> groups = m_upgradeHelper.createSequence(upgrade, context);
 
@@ -573,9 +577,8 @@ public class UpgradeHelperTest {
 
     Cluster cluster = makeCluster();
 
-    UpgradeContext context = m_upgradeContextFactory.create(cluster, UpgradeType.ROLLING,
-        Direction.DOWNGRADE, DOWNGRADE_VERSION, new HashMap<String, Object>());
-    context.setResolver(m_masterHostResolver);
+    UpgradeContext context = getMockUpgradeContext(cluster, Direction.DOWNGRADE,
+        UpgradeType.ROLLING, repositoryVersion2200);
 
     List<UpgradeGroupHolder> groups = m_upgradeHelper.createSequence(upgrade, context);
 
@@ -614,9 +617,8 @@ public class UpgradeHelperTest {
 
     Cluster cluster = makeCluster();
 
-    UpgradeContext context = m_upgradeContextFactory.create(cluster, UpgradeType.ROLLING,
-        Direction.UPGRADE, UPGRADE_VERSION, new HashMap<String, Object>());
-    context.setResolver(m_masterHostResolver);
+    UpgradeContext context = getMockUpgradeContext(cluster, Direction.UPGRADE, UpgradeType.ROLLING,
+        repositoryVersion2210);
 
     List<UpgradeGroupHolder> groups = m_upgradeHelper.createSequence(upgrade, context);
 
@@ -645,9 +647,8 @@ public class UpgradeHelperTest {
 
     Cluster cluster = makeCluster();
 
-    UpgradeContext context = m_upgradeContextFactory.create(cluster, UpgradeType.ROLLING,
-        Direction.UPGRADE, UPGRADE_VERSION, new HashMap<String, Object>());
-    context.setResolver(m_masterHostResolver);
+    UpgradeContext context = getMockUpgradeContext(cluster, Direction.UPGRADE, UpgradeType.ROLLING,
+        repositoryVersion2210);
 
     List<UpgradeGroupHolder> groups = m_upgradeHelper.createSequence(upgrade, context);
 
@@ -674,9 +675,8 @@ public class UpgradeHelperTest {
 
     Cluster cluster = makeCluster();
 
-    UpgradeContext context = m_upgradeContextFactory.create(cluster, UpgradeType.ROLLING,
-        Direction.UPGRADE, UPGRADE_VERSION, new HashMap<String, Object>());
-    context.setResolver(m_masterHostResolver);
+    UpgradeContext context = getMockUpgradeContext(cluster, Direction.UPGRADE, UpgradeType.ROLLING,
+        repositoryVersion2210);
 
     List<UpgradeGroupHolder> groups = m_upgradeHelper.createSequence(upgrade, context);
 
@@ -746,9 +746,8 @@ public class UpgradeHelperTest {
 
     Cluster cluster = makeCluster();
 
-    UpgradeContext context = m_upgradeContextFactory.create(cluster, UpgradeType.ROLLING,
-        Direction.UPGRADE, UPGRADE_VERSION, new HashMap<String, Object>());
-    context.setResolver(m_masterHostResolver);
+    UpgradeContext context = getMockUpgradeContext(cluster, Direction.UPGRADE, UpgradeType.ROLLING,
+        repositoryVersion2210);
 
     List<UpgradeGroupHolder> groups = m_upgradeHelper.createSequence(upgrade,
         context);
@@ -858,9 +857,8 @@ public class UpgradeHelperTest {
 
     Cluster cluster = makeCluster();
 
-    UpgradeContext context = m_upgradeContextFactory.create(cluster, UpgradeType.ROLLING,
-        Direction.UPGRADE, UPGRADE_VERSION, new HashMap<String, Object>());
-    context.setResolver(m_masterHostResolver);
+    UpgradeContext context = getMockUpgradeContext(cluster, Direction.UPGRADE, UpgradeType.ROLLING,
+        repositoryVersion2210);
 
     List<UpgradeGroupHolder> groups = m_upgradeHelper.createSequence(upgrade,
         context);
@@ -921,9 +919,8 @@ public class UpgradeHelperTest {
 
     Cluster cluster = makeCluster();
 
-    UpgradeContext context = m_upgradeContextFactory.create(cluster, UpgradeType.ROLLING,
-        Direction.UPGRADE, UPGRADE_VERSION, new HashMap<String, Object>());
-    context.setResolver(m_masterHostResolver);
+    UpgradeContext context = getMockUpgradeContext(cluster, Direction.UPGRADE, UpgradeType.ROLLING,
+        repositoryVersion2210);
 
     List<UpgradeGroupHolder> groups = m_upgradeHelper.createSequence(upgrade,
         context);
@@ -988,9 +985,8 @@ public class UpgradeHelperTest {
     assertNotNull(upgrade);
     Cluster cluster = makeCluster();
 
-    UpgradeContext context = m_upgradeContextFactory.create(cluster, UpgradeType.ROLLING,
-        Direction.UPGRADE, UPGRADE_VERSION, new HashMap<String, Object>());
-    context.setResolver(m_masterHostResolver);
+    UpgradeContext context = getMockUpgradeContext(cluster, Direction.UPGRADE, UpgradeType.ROLLING,
+        repositoryVersion2210);
 
     List<UpgradeGroupHolder> groups = m_upgradeHelper.createSequence(upgrade, context);
 
@@ -1065,9 +1061,8 @@ public class UpgradeHelperTest {
       numServiceChecksExpected++;
     }
 
-    UpgradeContext context = m_upgradeContextFactory.create(c, UpgradeType.ROLLING,
-        Direction.UPGRADE, UPGRADE_VERSION, new HashMap<String, Object>());
-    context.setResolver(m_masterHostResolver);
+    UpgradeContext context = getMockUpgradeContext(c, Direction.UPGRADE, UpgradeType.ROLLING,
+        repositoryVersion2210);
 
     List<UpgradeGroupHolder> groups = m_upgradeHelper.createSequence(upgrade, context);
 
@@ -1111,9 +1106,8 @@ public class UpgradeHelperTest {
 
     Cluster cluster = makeCluster();
 
-    UpgradeContext context = m_upgradeContextFactory.create(cluster, UpgradeType.ROLLING,
-        Direction.DOWNGRADE, DOWNGRADE_VERSION, new HashMap<String, Object>());
-    context.setResolver(m_masterHostResolver);
+    UpgradeContext context = getMockUpgradeContext(cluster, Direction.DOWNGRADE,
+        UpgradeType.ROLLING, repositoryVersion2200);
 
     List<UpgradeGroupHolder> groups = m_upgradeHelper.createSequence(upgrade, context);
 
@@ -1147,9 +1141,8 @@ public class UpgradeHelperTest {
 
     Cluster cluster = makeCluster();
 
-    UpgradeContext context = m_upgradeContextFactory.create(cluster, UpgradeType.ROLLING,
-        Direction.UPGRADE, UPGRADE_VERSION, new HashMap<String, Object>());
-    context.setResolver(m_masterHostResolver);
+    UpgradeContext context = getMockUpgradeContext(cluster, Direction.UPGRADE, UpgradeType.ROLLING,
+        repositoryVersion2210);
 
     List<UpgradeGroupHolder> groups = m_upgradeHelper.createSequence(upgrade, context);
 
@@ -1398,9 +1391,8 @@ public class UpgradeHelperTest {
 
     Cluster cluster = makeCluster();
 
-    UpgradeContext context = m_upgradeContextFactory.create(cluster, UpgradeType.ROLLING,
-        Direction.UPGRADE, UPGRADE_VERSION, new HashMap<String, Object>());
-    context.setResolver(m_masterHostResolver);
+    UpgradeContext context = getMockUpgradeContext(cluster, Direction.UPGRADE, UpgradeType.ROLLING,
+        repositoryVersion2210);
 
     List<UpgradeGroupHolder> groups = m_upgradeHelper.createSequence(upgrade, context);
 
@@ -1498,9 +1490,8 @@ public class UpgradeHelperTest {
     expect(m_masterHostResolver.getCluster()).andReturn(c).anyTimes();
     replay(m_masterHostResolver);
 
-    UpgradeContext context = m_upgradeContextFactory.create(c, UpgradeType.ROLLING,
-        Direction.DOWNGRADE, DOWNGRADE_VERSION, new HashMap<String, Object>());
-    context.setResolver(m_masterHostResolver);
+    UpgradeContext context = getMockUpgradeContext(c, Direction.DOWNGRADE, UpgradeType.ROLLING,
+        repositoryVersion2200);
 
     Map<String, UpgradePack> upgrades = ambariMetaInfo.getUpgradePacks("HDP", "2.1.1");
     assertTrue(upgrades.containsKey("upgrade_direction"));
@@ -1747,13 +1738,12 @@ public class UpgradeHelperTest {
     assertEquals(upgradeType, upgradePack.getType());
 
     // get an upgrade
-    UpgradeContext context = m_upgradeContextFactory.create(cluster, UpgradeType.ROLLING,
-        Direction.UPGRADE, UPGRADE_VERSION, new HashMap<String, Object>());
-
-    context.setResolver(m_masterHostResolver);
+    UpgradeContext context = getMockUpgradeContext(cluster, Direction.UPGRADE, UpgradeType.ROLLING,
+        repositoryVersion2210, m_masterHostResolver, false);
 
     context.setSupportedServices(Collections.singleton("ZOOKEEPER"));
     context.setScope(UpgradeScope.COMPLETE);
+    replay(context);
 
     List<Grouping> groupings = upgradePack.getGroups(Direction.UPGRADE);
     assertEquals(2, groupings.size());
@@ -1850,9 +1840,8 @@ public class UpgradeHelperTest {
 
     MasterHostResolver resolver = new MasterHostResolver(m_configHelper, c);
 
-    UpgradeContext context = m_upgradeContextFactory.create(c, UpgradeType.NON_ROLLING,
-        Direction.UPGRADE, "2.2.0", new HashMap<String, Object>());
-    context.setResolver(resolver);
+    UpgradeContext context = getMockUpgradeContext(c, Direction.UPGRADE, UpgradeType.NON_ROLLING,
+        repositoryVersion2200, resolver, true);
 
     List<UpgradeGroupHolder> groups = m_upgradeHelper.createSequence(upgradePack, context);
 
@@ -1862,9 +1851,8 @@ public class UpgradeHelperTest {
     sch2.setVersion("2.1.1");
     resolver = new MasterHostResolver(m_configHelper, c, "2.1.1");
 
-    context = m_upgradeContextFactory.create(c, UpgradeType.NON_ROLLING, Direction.DOWNGRADE,
-        "2.1.1", new HashMap<String, Object>());
-    context.setResolver(resolver);
+    context = getMockUpgradeContext(c, Direction.DOWNGRADE, UpgradeType.NON_ROLLING,
+        repositoryVersion2110, resolver, true);
 
     groups = m_upgradeHelper.createSequence(upgradePack, context);
 
@@ -1947,9 +1935,8 @@ public class UpgradeHelperTest {
     UpgradePack upgrade = upgrades.get("upgrade_multi_server_tasks");
     assertNotNull(upgrade);
 
-    UpgradeContext context = m_upgradeContextFactory.create(c, UpgradeType.NON_ROLLING,
-        Direction.UPGRADE, "2.2.0", new HashMap<String, Object>());
-    context.setResolver(m_masterHostResolver);
+    UpgradeContext context = getMockUpgradeContext(c, Direction.UPGRADE, UpgradeType.NON_ROLLING,
+        repositoryVersion2200);
 
     List<UpgradeGroupHolder> groups = m_upgradeHelper.createSequence(upgrade, context);
 
@@ -2134,10 +2121,9 @@ public class UpgradeHelperTest {
 
     MasterHostResolver resolver = new MasterHostResolver(m_configHelper, c);
 
-    UpgradeContext context = m_upgradeContextFactory.create(c, UpgradeType.HOST_ORDERED,
-        Direction.UPGRADE, "2.2.0", new HashMap<String, Object>());
+    UpgradeContext context = getMockUpgradeContext(c, Direction.UPGRADE, UpgradeType.HOST_ORDERED,
+        repositoryVersion2200, resolver, true);
 
-    context.setResolver(resolver);
     List<UpgradeGroupHolder> groups = m_upgradeHelper.createSequence(upgradePack, context);
 
     assertEquals(1, groups.size());
@@ -2175,8 +2161,8 @@ public class UpgradeHelperTest {
     hbaseMaster1.setVersion("2.1.1");
     resolver = new MasterHostResolver(m_configHelper, c, "2.1.1");
 
-    m_upgradeContextFactory.create(c, UpgradeType.HOST_ORDERED, Direction.DOWNGRADE,
-        "2.1.1", new HashMap<String, Object>());
+    context = getMockUpgradeContext(c, Direction.DOWNGRADE, UpgradeType.HOST_ORDERED,
+        repositoryVersion2110, resolver, true);
 
     context.setResolver(resolver);
     groups = m_upgradeHelper.createSequence(upgradePack, context);
@@ -2190,10 +2176,9 @@ public class UpgradeHelperTest {
     hbaseMaster1.setVersion("2.1.1");
     resolver = new MasterHostResolver(m_configHelper, c, "2.1.1");
 
-    m_upgradeContextFactory.create(c, UpgradeType.HOST_ORDERED, Direction.DOWNGRADE,
-        "2.1.1", new HashMap<String, Object>());
+    context = getMockUpgradeContext(c, Direction.DOWNGRADE, UpgradeType.HOST_ORDERED,
+        repositoryVersion2110, resolver, true);
 
-    context.setResolver(resolver);
     groups = m_upgradeHelper.createSequence(upgradePack, context);
 
     assertEquals(1, groups.size());
@@ -2215,10 +2200,8 @@ public class UpgradeHelperTest {
 
     Cluster cluster = makeCluster();
 
-    UpgradeContext context = m_upgradeContextFactory.create(cluster, UpgradeType.ROLLING,
-        Direction.UPGRADE, UPGRADE_VERSION, new HashMap<String, Object>());
-
-    context.setResolver(m_masterHostResolver);
+    UpgradeContext context = getMockUpgradeContext(cluster, Direction.UPGRADE, UpgradeType.ROLLING,
+        repositoryVersion2210);
 
     // initially, no conditions should be met
     List<UpgradeGroupHolder> groups = m_upgradeHelper.createSequence(upgrade, context);
@@ -2252,6 +2235,62 @@ public class UpgradeHelperTest {
   }
 
   /**
+   * @param cluster
+   * @param direction
+   * @param type
+   * @return
+   */
+  private UpgradeContext getMockUpgradeContext(Cluster cluster, Direction direction,
+      UpgradeType type) {
+    return getMockUpgradeContext(cluster, direction, type, repositoryVersion2210);
+  }
+
+  /**
+   * @param cluster
+   * @param direction
+   * @param type
+   * @return
+   */
+  private UpgradeContext getMockUpgradeContext(Cluster cluster, Direction direction,
+      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, MasterHostResolver resolver,
+      boolean replay) {
+    RepositoryVersionEntity sourceRepositoryVersion = cluster.getAllClusterVersions().iterator().next().getRepositoryVersion();
+
+    UpgradeContext context = EasyMock.createNiceMock(UpgradeContext.class);
+    expect(context.getCluster()).andReturn(cluster).anyTimes();
+    expect(context.getType()).andReturn(type).anyTimes();
+    expect(context.getDirection()).andReturn(direction).anyTimes();
+    expect(context.getSourceRepositoryVersion()).andReturn(sourceRepositoryVersion).anyTimes();
+    expect(context.getTargetRepositoryVersion()).andReturn(repositoryVersion).anyTimes();
+    expect(context.getAmbariMetaInfo()).andReturn(ambariMetaInfo).anyTimes();
+
+    // only set this if supplied
+    if (null != resolver) {
+      expect(context.getResolver()).andReturn(resolver).anyTimes();
+    }
+
+    if (replay) {
+      replay(context);
+    }
+
+    return context;
+  }
+
+  /**
    * Extend {@link org.apache.ambari.server.stack.MasterHostResolver} in order
    * to overwrite the JMX methods.
    */

http://git-wip-us.apache.org/repos/asf/ambari/blob/7ad307c2/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 a61b7b1..22183ba 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
@@ -142,18 +142,21 @@ public class ClusterEffectiveVersionTest extends EasyMockSupport {
 
     Cluster clusterSpy = Mockito.spy(m_cluster);
 
+    RepositoryVersionEntity repoVersion2300 = createNiceMock(RepositoryVersionEntity.class);
+    RepositoryVersionEntity repoVersion2400 = createNiceMock(RepositoryVersionEntity.class);
+
+    EasyMock.expect(repoVersion2300.getVersion()).andReturn("2.3.0.0-1234").anyTimes();
+    EasyMock.expect(repoVersion2400.getVersion()).andReturn("2.4.0.0-1234").anyTimes();
+
     UpgradeEntity upgradeEntity = createNiceMock(UpgradeEntity.class);
     EasyMock.expect(upgradeEntity.getId()).andReturn(1L).atLeastOnce();
     EasyMock.expect(upgradeEntity.getUpgradeType()).andReturn(UpgradeType.ROLLING).atLeastOnce();
-    EasyMock.expect(upgradeEntity.getFromVersion()).andReturn("2.3.0.0-1234").anyTimes();
-    EasyMock.expect(upgradeEntity.getToVersion()).andReturn("2.4.0.0-1234").atLeastOnce();
-
-    RepositoryVersionEntity repositoryVersionEntity = createNiceMock(RepositoryVersionEntity.class);
-    EasyMock.expect(repositoryVersionEntity.getVersion()).andReturn("2.4.0.0-1234").atLeastOnce();
+    EasyMock.expect(upgradeEntity.getFromRepositoryVersion()).andReturn(repoVersion2300).anyTimes();
+    EasyMock.expect(upgradeEntity.getToRepositoryVersion()).andReturn(repoVersion2400).atLeastOnce();
 
     ClusterVersionEntity clusterVersionUpgradingTo = createNiceMock(ClusterVersionEntity.class);
     EasyMock.expect(clusterVersionUpgradingTo.getRepositoryVersion()).andReturn(
-        repositoryVersionEntity).atLeastOnce();
+        repoVersion2400).atLeastOnce();
 
     List<ClusterVersionEntity> clusterVersionEntities = Lists.newArrayList(clusterVersionUpgradingTo);
     EasyMock.expect(m_clusterEntity.getClusterVersionEntities()).andReturn(clusterVersionEntities).atLeastOnce();
@@ -183,20 +186,24 @@ public class ClusterEffectiveVersionTest extends EasyMockSupport {
 
     Cluster clusterSpy = Mockito.spy(m_cluster);
 
+    RepositoryVersionEntity repoVersion2300 = createNiceMock(RepositoryVersionEntity.class);
+    RepositoryVersionEntity repoVersion2400 = createNiceMock(RepositoryVersionEntity.class);
+
+    EasyMock.expect(repoVersion2300.getVersion()).andReturn("2.3.0.0-1234").anyTimes();
+    EasyMock.expect(repoVersion2400.getVersion()).andReturn("2.4.0.0-1234").anyTimes();
+
     // from/to are switched on downgrade
     UpgradeEntity upgradeEntity = createNiceMock(UpgradeEntity.class);
     EasyMock.expect(upgradeEntity.getId()).andReturn(1L).atLeastOnce();
     EasyMock.expect(upgradeEntity.getUpgradeType()).andReturn(UpgradeType.NON_ROLLING).atLeastOnce();
-    EasyMock.expect(upgradeEntity.getToVersion()).andReturn("2.3.0.0-1234").atLeastOnce();
-    EasyMock.expect(upgradeEntity.getFromVersion()).andReturn("2.4.0.0-1234").anyTimes();
-    EasyMock.expect(upgradeEntity.getDirection()).andReturn(Direction.DOWNGRADE).atLeastOnce();
 
-    RepositoryVersionEntity repositoryVersionEntity = createNiceMock(RepositoryVersionEntity.class);
-    EasyMock.expect(repositoryVersionEntity.getVersion()).andReturn("2.3.0.0-1234").atLeastOnce();
+    EasyMock.expect(upgradeEntity.getFromRepositoryVersion()).andReturn(repoVersion2400).anyTimes();
+    EasyMock.expect(upgradeEntity.getToRepositoryVersion()).andReturn(repoVersion2300).atLeastOnce();
+    EasyMock.expect(upgradeEntity.getDirection()).andReturn(Direction.DOWNGRADE).atLeastOnce();
 
     ClusterVersionEntity clusterVersionUpgradingTo = createNiceMock(ClusterVersionEntity.class);
     EasyMock.expect(clusterVersionUpgradingTo.getRepositoryVersion()).andReturn(
-        repositoryVersionEntity).atLeastOnce();
+        repoVersion2300).atLeastOnce();
 
     List<ClusterVersionEntity> clusterVersionEntities = Lists.newArrayList(clusterVersionUpgradingTo);
     EasyMock.expect(m_clusterEntity.getClusterVersionEntities()).andReturn(clusterVersionEntities).atLeastOnce();

http://git-wip-us.apache.org/repos/asf/ambari/blob/7ad307c2/ambari-server/src/test/java/org/apache/ambari/server/state/services/RetryUpgradeActionServiceTest.java
----------------------------------------------------------------------
diff --git a/ambari-server/src/test/java/org/apache/ambari/server/state/services/RetryUpgradeActionServiceTest.java b/ambari-server/src/test/java/org/apache/ambari/server/state/services/RetryUpgradeActionServiceTest.java
index 9570054..85c08fe 100644
--- a/ambari-server/src/test/java/org/apache/ambari/server/state/services/RetryUpgradeActionServiceTest.java
+++ b/ambari-server/src/test/java/org/apache/ambari/server/state/services/RetryUpgradeActionServiceTest.java
@@ -81,6 +81,8 @@ public class RetryUpgradeActionServiceTest {
   Long upgradeRequestId = 1L;
   Long stageId = 1L;
 
+  RepositoryVersionEntity repoVersion2200;
+
   @Before
   public void before() throws NoSuchFieldException, IllegalAccessException {
     injector = Guice.createInjector(new InMemoryDefaultTestModule());
@@ -241,12 +243,12 @@ public class RetryUpgradeActionServiceTest {
     clusters.addCluster(clusterName, stack220);
     cluster = clusters.getCluster("c1");
 
-    RepositoryVersionEntity repoVersionEntity = new RepositoryVersionEntity();
-    repoVersionEntity.setDisplayName("Initial Version");
-    repoVersionEntity.setOperatingSystems("");
-    repoVersionEntity.setStack(stackEntity220);
-    repoVersionEntity.setVersion("2.2.0.0");
-    repoVersionDAO.create(repoVersionEntity);
+    repoVersion2200 = new RepositoryVersionEntity();
+    repoVersion2200.setDisplayName("Initial Version");
+    repoVersion2200.setOperatingSystems("");
+    repoVersion2200.setStack(stackEntity220);
+    repoVersion2200.setVersion("2.2.0.0");
+    repoVersionDAO.create(repoVersion2200);
 
     helper.getOrCreateRepositoryVersion(stack220, stack220.getStackVersion());
 
@@ -259,12 +261,12 @@ public class RetryUpgradeActionServiceTest {
    * @throws AmbariException
    */
   private void prepareUpgrade() throws AmbariException {
-    RepositoryVersionEntity repoVersionEntity = new RepositoryVersionEntity();
-    repoVersionEntity.setDisplayName("Version to Upgrade To");
-    repoVersionEntity.setOperatingSystems("");
-    repoVersionEntity.setStack(stackEntity220);
-    repoVersionEntity.setVersion("2.2.0.1");
-    repoVersionDAO.create(repoVersionEntity);
+    RepositoryVersionEntity repoVersion2201 = new RepositoryVersionEntity();
+    repoVersion2201.setDisplayName("Version to Upgrade To");
+    repoVersion2201.setOperatingSystems("");
+    repoVersion2201.setStack(stackEntity220);
+    repoVersion2201.setVersion("2.2.0.1");
+    repoVersionDAO.create(repoVersion2201);
 
     helper.getOrCreateRepositoryVersion(stack220, stack220.getStackVersion());
 
@@ -291,8 +293,8 @@ public class RetryUpgradeActionServiceTest {
     upgrade.setUpgradePackage("some-name");
     upgrade.setUpgradeType(UpgradeType.ROLLING);
     upgrade.setDirection(Direction.UPGRADE);
-    upgrade.setFromVersion("2.2.0.0");
-    upgrade.setToVersion("2.2.0.1");
+    upgrade.setFromRepositoryVersion(repoVersion2200);
+    upgrade.setToRepositoryVersion(repoVersion2201);
     upgradeDAO.create(upgrade);
 
     cluster.setUpgradeEntity(upgrade);