AMBARI-13491 - Expose Failure Auto-Skip Setting On Upgrade Endpoint (jonathanhurley)

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

@@ -465,6 +465,8 @@ public class UpgradeResourceProvider extends AbstractControllerResourceProvider
     setResourceProperty(resource, UPGRADE_TO_VERSION, entity.getToVersion(), requestedIds);
     setResourceProperty(resource, UPGRADE_DIRECTION, entity.getDirection(), requestedIds);
     setResourceProperty(resource, UPGRADE_DOWNGRADE_ALLOWED, entity.isDowngradeAllowed(), requestedIds);
+    setResourceProperty(resource, UPGRADE_SKIP_FAILURES, entity.isComponentFailureAutoSkipped(), requestedIds);
+    setResourceProperty(resource, UPGRADE_SKIP_SC_FAILURES, entity.isServiceCheckFailureAutoSkipped(), requestedIds);
 
     return resource;
   }
@@ -762,7 +764,7 @@ public class UpgradeResourceProvider extends AbstractControllerResourceProvider
               itemEntity.setTasks(wrapper.getTasksJson());
               itemEntity.setHosts(wrapper.getHostsJson());
               itemEntities.add(itemEntity);
-              
+
               // At this point, need to change the effective Stack Id so that subsequent tasks run on the newer value.
               // TODO AMBARI-12698, check if this works during a Stop-the-World Downgrade.
               if (UpdateStackGrouping.class.equals(group.groupClass)) {
@@ -806,6 +808,8 @@ public class UpgradeResourceProvider extends AbstractControllerResourceProvider
     entity.setDirection(direction);
     entity.setUpgradePackage(pack.getName());
     entity.setUpgradeType(pack.getType());
+    entity.setAutoSkipComponentFailures(skipComponentFailures);
+    entity.setAutoSkipServiceCheckFailures(skipServiceCheckFailures);
     entity.setDowngradeAllowed(pack.isDowngradeAllowed());
 
     req.getRequestStatusResponse();

ambari-server/src/main/java/org/apache/ambari/server/orm/entities/UpgradeEntity.java

@@ -33,14 +33,15 @@ import javax.persistence.OneToMany;
 import javax.persistence.Table;
 import javax.persistence.TableGenerator;
 
+import org.apache.ambari.server.actionmanager.HostRoleStatus;
 import org.apache.ambari.server.state.stack.upgrade.Direction;
 import org.apache.ambari.server.state.stack.upgrade.UpgradeType;
 
 /**
  * Models the data representation of an upgrade
  */
-@Table(name = "upgrade")
 @Entity
+@Table(name = "upgrade")
 @TableGenerator(name = "upgrade_id_generator",
     table = "ambari_sequences", pkColumnName = "sequence_name", valueColumnName = "sequence_value",
     pkColumnValue = "upgrade_id_seq", initialValue = 0)
@@ -84,6 +85,12 @@ public class UpgradeEntity {
   @Enumerated(value = EnumType.STRING)
   private UpgradeType upgradeType;
 
+  @Column(name = "skip_failures", nullable = false)
+  private Integer skipFailures = 0;
+
+  @Column(name = "skip_sc_failures", nullable = false)
+  private Integer skipServiceCheckFailures = 0;
+
   @Column(name="downgrade_allowed", nullable = false)
   private Integer downgrade_allowed = 1;
 
@@ -207,7 +214,7 @@ public class UpgradeEntity {
   /**
    * @param canDowngrade {@code true} to allow downgrade, {@code false} to disallow downgrade
    */
-  public void setDowngradeAllowed(boolean canDowngrade){ this.downgrade_allowed = (!canDowngrade ? 0 : 1); }
+  public void setDowngradeAllowed(boolean canDowngrade){ downgrade_allowed = (!canDowngrade ? 0 : 1); }
 
   /**
    * @param upgradeType the upgrade type to set
@@ -230,6 +237,52 @@ public class UpgradeEntity {
     this.upgradePackage = upgradePackage;
   }
 
+  /**
+   * Gets whether skippable components that failed are automatically skipped.
+   * They will be placed into the {@link HostRoleStatus#SKIPPED_FAILED} state.
+   *
+   * @return {@code true} if skippable failed components are automatically
+   *         skipped when they fail.
+   */
+  public boolean isComponentFailureAutoSkipped() {
+    return skipFailures != 0;
+  }
+
+  /**
+   * Sets whether skippable components that failed are automatically skipped.
+   *
+   * @param autoSkipComponentFailures
+   *          {@code true} to automatically skip component failures which are
+   *          marked as skippable.
+   */
+  public void setAutoSkipComponentFailures(boolean autoSkipComponentFailures) {
+    skipFailures = autoSkipComponentFailures ? 1 : 0;
+  }
+
+  /**
+   * Gets whether skippable service checks that failed are automatically
+   * skipped. They will be placed into the {@link HostRoleStatus#SKIPPED_FAILED}
+   * state.
+   *
+   * @return {@code true} if service checks are automatically skipped when they
+   *         fail.
+   */
+  public boolean isServiceCheckFailureAutoSkipped() {
+    return skipServiceCheckFailures != 0;
+  }
+
+  /**
+   * Sets whether skippable service checks that failed are automatically
+   * skipped.
+   *
+   * @param autoSkipServiceCheckFailures
+   *          {@code true} to automatically skip service check failures which
+   *          are marked as being skippable.
+   */
+  public void setAutoSkipServiceCheckFailures(boolean autoSkipServiceCheckFailures) {
+    skipServiceCheckFailures = autoSkipServiceCheckFailures ? 1 : 0;
+  }
+
   @Override
   public boolean equals(Object o) {
     if (this == o) {

ambari-server/src/main/java/org/apache/ambari/server/upgrade/UpgradeCatalog213.java

@@ -18,10 +18,16 @@
 
 package org.apache.ambari.server.upgrade;
 
-import com.google.gson.JsonObject;
-import com.google.gson.JsonParser;
-import com.google.inject.Inject;
-import com.google.inject.Injector;
+import java.sql.SQLException;
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Set;
+import java.util.UUID;
+
 import org.apache.ambari.server.AmbariException;
 import org.apache.ambari.server.controller.AmbariManagementController;
 import org.apache.ambari.server.orm.DBAccessor;
@@ -40,15 +46,10 @@ import org.apache.ambari.server.utils.VersionUtils;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
-import java.sql.SQLException;
-import java.util.ArrayList;
-import java.util.Collections;
-import java.util.HashMap;
-import java.util.HashSet;
-import java.util.List;
-import java.util.Map;
-import java.util.Set;
-import java.util.UUID;
+import com.google.gson.JsonObject;
+import com.google.gson.JsonParser;
+import com.google.inject.Inject;
+import com.google.inject.Injector;
 
 /**
  * Upgrade catalog for version 2.1.3.
@@ -73,7 +74,11 @@ public class UpgradeCatalog213 extends AbstractUpgradeCatalog {
                                     "  ulimit -l {{datanode_max_locked_memory}}\n" +
                                     "fi\n" +
                                     "{% endif %};\n";
+
   private static final String DOWNGRADE_ALLOWED_COLUMN = "downgrade_allowed";
+  private static final String UPGRADE_SKIP_FAILURE_COLUMN = "skip_failures";
+  private static final String UPGRADE_SKIP_SC_FAILURE_COLUMN = "skip_sc_failures";
+
   private static final String KERBEROS_DESCRIPTOR_TABLE = "kerberos_descriptor";
   private static final String KERBEROS_DESCRIPTOR_NAME_COLUMN = "kerberos_descriptor_name";
   private static final String KERBEROS_DESCRIPTOR_COLUMN = "kerberos_descriptor";
@@ -131,11 +136,10 @@ public class UpgradeCatalog213 extends AbstractUpgradeCatalog {
   }
 
   protected void executeUpgradeDDLUpdates() throws AmbariException, SQLException {
-    dbAccessor.addColumn(UPGRADE_TABLE, new DBColumnInfo(DOWNGRADE_ALLOWED_COLUMN, Short.class, 1, null, true));
+    updateUpgradesDDL();
   }
 
   private void addKerberosDescriptorTable() throws SQLException {
-
     List<DBAccessor.DBColumnInfo> columns = new ArrayList<DBAccessor.DBColumnInfo>();
     columns.add(new DBAccessor.DBColumnInfo(KERBEROS_DESCRIPTOR_NAME_COLUMN, String.class, 255, null, false));
     columns.add(new DBAccessor.DBColumnInfo(KERBEROS_DESCRIPTOR_COLUMN, char[].class, null, null, false));
@@ -149,21 +153,43 @@ public class UpgradeCatalog213 extends AbstractUpgradeCatalog {
    */
   @Override
   protected void executePreDMLUpdates() throws AmbariException, SQLException {
-    populateDowngradeAllowed();
+    executeUpgradePreDMLUpdates();
   }
 
-  protected void populateDowngradeAllowed() throws AmbariException, SQLException {
+  /**
+   * Updates the following columns on the {@value #UPGRADE_TABLE} table to
+   * default values:
+   * <ul>
+   * <li>{value {@link #DOWNGRADE_ALLOWED_COLUMN}}</li>
+   * <li>{value {@link #UPGRADE_SKIP_FAILURE_COLUMN}}</li>
+   * <li>{value {@link #UPGRADE_SKIP_SC_FAILURE_COLUMN}}</li>
+   * </ul>
+   *
+   * @throws AmbariException
+   * @throws SQLException
+   */
+  protected void executeUpgradePreDMLUpdates() throws AmbariException, SQLException {
     UpgradeDAO upgradeDAO = injector.getInstance(UpgradeDAO.class);
     List<UpgradeEntity> upgrades = upgradeDAO.findAll();
     for (UpgradeEntity upgrade: upgrades){
       if (upgrade.isDowngradeAllowed() == null) {
         upgrade.setDowngradeAllowed(true);
         upgradeDAO.merge(upgrade);
-        LOG.info(String.format("Update upgrade id %s, upgrade pack %s from version %s to %s",
-          upgrade.getId(), upgrade.getUpgradePackage(), upgrade.getFromVersion(), upgrade.getToVersion()));
       }
+
+      // ensure that these are set to false for existing upgrades
+      upgrade.setAutoSkipComponentFailures(false);
+      upgrade.setAutoSkipServiceCheckFailures(false);
+
+      LOG.info(String.format("Updated upgrade id %s, upgrade pack %s from version %s to %s",
+          upgrade.getId(), upgrade.getUpgradePackage(), upgrade.getFromVersion(),
+          upgrade.getToVersion()));
     }
+
+    // make the columns nullable now that they have defaults
     dbAccessor.setColumnNullable(UPGRADE_TABLE, DOWNGRADE_ALLOWED_COLUMN, false);
+    dbAccessor.setColumnNullable(UPGRADE_TABLE, UPGRADE_SKIP_FAILURE_COLUMN, false);
+    dbAccessor.setColumnNullable(UPGRADE_TABLE, UPGRADE_SKIP_SC_FAILURE_COLUMN, false);
   }
 
   @Override
@@ -179,6 +205,22 @@ public class UpgradeCatalog213 extends AbstractUpgradeCatalog {
     updateZookeeperLog4j();
   }
 
+  /**
+   * Adds the following columns to the {@value #UPGRADE_TABLE} table:
+   * <ul>
+   * <li>{@value #DOWNGRADE_ALLOWED_COLUMN}</li>
+   * <li>{@value #UPGRADE_SKIP_FAILURE_COLUMN}</li>
+   * <li>{@value #UPGRADE_SKIP_SC_FAILURE_COLUMN}</li>
+   * </ul>
+   *
+   * @throws SQLException
+   */
+  protected void updateUpgradesDDL() throws SQLException{
+    dbAccessor.addColumn(UPGRADE_TABLE, new DBColumnInfo(DOWNGRADE_ALLOWED_COLUMN, Short.class, 1, null, true));
+    dbAccessor.addColumn(UPGRADE_TABLE, new DBColumnInfo(UPGRADE_SKIP_FAILURE_COLUMN, Short.class, 1, null, true));
+    dbAccessor.addColumn(UPGRADE_TABLE, new DBColumnInfo(UPGRADE_SKIP_SC_FAILURE_COLUMN, Short.class, 1, null, true));
+  }
+
   /**
    * Modifies the JSON of some of the alert definitions which have changed
    * between Ambari versions.
@@ -247,7 +289,7 @@ public class UpgradeCatalog213 extends AbstractUpgradeCatalog {
 
     return rootJson.toString();
   }
-  
+
   protected void updateHadoopEnv() throws AmbariException {
     AmbariManagementController ambariManagementController = injector.getInstance(AmbariManagementController.class);
 

ambari-server/src/main/resources/Ambari-DDL-MySQL-CREATE.sql

@@ -891,6 +891,8 @@ CREATE TABLE upgrade (
   direction VARCHAR(255) DEFAULT 'UPGRADE' NOT NULL,
   upgrade_package VARCHAR(255) NOT NULL,
   upgrade_type VARCHAR(32) NOT NULL,
+  skip_failures TINYINT(1) NOT NULL DEFAULT 0,
+  skip_sc_failures TINYINT(1) NOT NULL DEFAULT 0,
   downgrade_allowed TINYINT(1) NOT NULL DEFAULT 1,
   PRIMARY KEY (upgrade_id),
   FOREIGN KEY (cluster_id) REFERENCES clusters(cluster_id),

ambari-server/src/main/resources/Ambari-DDL-Oracle-CREATE.sql

@@ -880,6 +880,8 @@ CREATE TABLE upgrade (
   direction VARCHAR2(255) DEFAULT 'UPGRADE' NOT NULL,
   upgrade_package VARCHAR2(255) NOT NULL,
   upgrade_type VARCHAR2(32) NOT NULL,
+  skip_failures NUMBER(1) DEFAULT 0 NOT NULL,
+  skip_sc_failures NUMBER(1) DEFAULT 0 NOT NULL,
   downgrade_allowed NUMBER(1) DEFAULT 1 NOT NULL,
   PRIMARY KEY (upgrade_id),
   FOREIGN KEY (cluster_id) REFERENCES clusters(cluster_id),

ambari-server/src/main/resources/Ambari-DDL-Postgres-CREATE.sql

@@ -882,6 +882,8 @@ CREATE TABLE upgrade (
   direction VARCHAR(255) DEFAULT 'UPGRADE' NOT NULL,
   upgrade_package VARCHAR(255) NOT NULL,
   upgrade_type VARCHAR(32) NOT NULL,
+  skip_failures SMALLINT DEFAULT 0 NOT NULL,
+  skip_sc_failures SMALLINT DEFAULT 0 NOT NULL,
   downgrade_allowed SMALLINT DEFAULT 1 NOT NULL,
   PRIMARY KEY (upgrade_id),
   FOREIGN KEY (cluster_id) REFERENCES clusters(cluster_id),

ambari-server/src/main/resources/Ambari-DDL-Postgres-EMBEDDED-CREATE.sql

@@ -974,6 +974,8 @@ CREATE TABLE ambari.upgrade (
   direction VARCHAR(255) DEFAULT 'UPGRADE' NOT NULL,
   upgrade_package VARCHAR(255) NOT NULL,
   upgrade_type VARCHAR(32) NOT NULL,
+  skip_failures SMALLINT DEFAULT 0 NOT NULL,
+  skip_sc_failures SMALLINT DEFAULT 0 NOT NULL,
   downgrade_allowed SMALLINT DEFAULT 1 NOT NULL,
   PRIMARY KEY (upgrade_id),
   FOREIGN KEY (cluster_id) REFERENCES ambari.clusters(cluster_id),

ambari-server/src/main/resources/Ambari-DDL-SQLAnywhere-CREATE.sql

@@ -879,6 +879,8 @@ CREATE TABLE upgrade (
   from_version VARCHAR(255) DEFAULT '' NOT NULL,
   to_version VARCHAR(255) DEFAULT '' NOT NULL,
   direction VARCHAR(255) DEFAULT 'UPGRADE' NOT NULL,
+  skip_failures BIT NOT NULL DEFAULT 0,
+  skip_sc_failures BIT NOT NULL DEFAULT 0,  
   downgrade_allowed BIT NOT NULL DEFAULT 1,
   PRIMARY KEY (upgrade_id),
   FOREIGN KEY (cluster_id) REFERENCES clusters(cluster_id),

ambari-server/src/main/resources/Ambari-DDL-SQLServer-CREATE.sql

@@ -997,6 +997,8 @@ CREATE TABLE upgrade (
   direction VARCHAR(255) DEFAULT 'UPGRADE' NOT NULL,
   upgrade_package VARCHAR(255) NOT NULL,
   upgrade_type VARCHAR(32) NOT NULL,
+  skip_failures BIT NOT NULL DEFAULT 0,
+  skip_sc_failures BIT NOT NULL DEFAULT 0,
   downgrade_allowed BIT NOT NULL DEFAULT 1,
   PRIMARY KEY CLUSTERED (upgrade_id),
   FOREIGN KEY (cluster_id) REFERENCES clusters(cluster_id),

ambari-server/src/test/java/org/apache/ambari/server/controller/internal/UpgradeResourceProviderTest.java

@@ -415,8 +415,10 @@ public class UpgradeResourceProviderTest {
     assertEquals(1, resources.size());
     res = resources.iterator().next();
     assertNotNull(res.getPropertyValue("Upgrade/progress_percent"));
-    assertNotNull(res.getPropertyValue("Upgrade/direction"));
-    assertEquals(Direction.UPGRADE, res.getPropertyValue("Upgrade/direction"));
+    assertNotNull(res.getPropertyValue(UpgradeResourceProvider.UPGRADE_DIRECTION));
+    assertEquals(Direction.UPGRADE, res.getPropertyValue(UpgradeResourceProvider.UPGRADE_DIRECTION));
+    assertEquals(false, res.getPropertyValue(UpgradeResourceProvider.UPGRADE_SKIP_FAILURES));
+    assertEquals(false, res.getPropertyValue(UpgradeResourceProvider.UPGRADE_SKIP_SC_FAILURES));
 
     // upgrade groups
     propertyIds.clear();
@@ -478,6 +480,52 @@ public class UpgradeResourceProviderTest {
     assertTrue(res.getPropertyValue("UpgradeItem/text").toString().startsWith("Please confirm"));
   }
 
+  /**
+   * Tests that retrieving an upgrade correctly populates less common upgrade
+   * options correctly.
+   */
+  @Test
+  public void testGetResourcesWithSpecialOptions() throws Exception {
+    Cluster cluster = clusters.getCluster("c1");
+
+    List<UpgradeEntity> upgrades = upgradeDao.findUpgrades(cluster.getClusterId());
+    assertEquals(0, upgrades.size());
+
+    Map<String, Object> requestProps = new HashMap<String, Object>();
+    requestProps.put(UpgradeResourceProvider.UPGRADE_CLUSTER_NAME, "c1");
+    requestProps.put(UpgradeResourceProvider.UPGRADE_VERSION, "2.1.1.1");
+    requestProps.put(UpgradeResourceProvider.UPGRADE_PACK, "upgrade_test");
+    requestProps.put(UpgradeResourceProvider.UPGRADE_SKIP_PREREQUISITE_CHECKS, "true");
+
+    // tests skipping SC failure options
+    requestProps.put(UpgradeResourceProvider.UPGRADE_SKIP_FAILURES, "true");
+    requestProps.put(UpgradeResourceProvider.UPGRADE_SKIP_SC_FAILURES, "true");
+
+    ResourceProvider upgradeResourceProvider = createProvider(amc);
+
+    Request request = PropertyHelper.getCreateRequest(Collections.singleton(requestProps), null);
+    RequestStatus status = upgradeResourceProvider.createResources(request);
+    assertNotNull(status);
+
+    // upgrade
+    Set<String> propertyIds = new HashSet<String>();
+    propertyIds.add("Upgrade");
+
+    Predicate predicate = new PredicateBuilder()
+      .property(UpgradeResourceProvider.UPGRADE_REQUEST_ID).equals("1").and()
+      .property(UpgradeResourceProvider.UPGRADE_CLUSTER_NAME).equals("c1")
+      .toPredicate();
+
+    request = PropertyHelper.getReadRequest(propertyIds);
+    Set<Resource> resources = upgradeResourceProvider.getResources(request, predicate);
+
+    assertEquals(1, resources.size());
+    Resource resource = resources.iterator().next();
+
+    assertEquals(true, resource.getPropertyValue(UpgradeResourceProvider.UPGRADE_SKIP_FAILURES));
+    assertEquals(true, resource.getPropertyValue(UpgradeResourceProvider.UPGRADE_SKIP_SC_FAILURES));
+  }
+
   @Ignore
   @Test
   public void testCreatePartialDowngrade() throws Exception {

ambari-server/src/test/java/org/apache/ambari/server/orm/dao/UpgradeDAOTest.java

@@ -177,4 +177,35 @@ public class UpgradeDAOTest {
     assertNotNull(lastUpgradeForCluster);
     assertEquals(33L, (long)lastUpgradeForCluster.getId());
   }
+
+  /**
+   * Tests that certain columns in an {@link UpgradeEntity} are updatable.
+   *
+   * @throws Exception
+   */
+  @Test
+  public void testUpdatableColumns() throws Exception {
+    UpgradeEntity upgradeEntity = new UpgradeEntity();
+    upgradeEntity.setId(11L);
+    upgradeEntity.setClusterId(Long.valueOf(1));
+    upgradeEntity.setDirection(Direction.UPGRADE);
+    upgradeEntity.setRequestId(Long.valueOf(1));
+    upgradeEntity.setFromVersion("2.2.0.0-1234");
+    upgradeEntity.setToVersion("2.3.0.0-4567");
+    upgradeEntity.setUpgradeType(UpgradeType.ROLLING);
+    upgradeEntity.setUpgradePackage("test-upgrade");
+    dao.create(upgradeEntity);
+
+    UpgradeEntity lastUpgradeForCluster = dao.findLastUpgradeForCluster(1);
+    Assert.assertFalse(lastUpgradeForCluster.isComponentFailureAutoSkipped());
+    Assert.assertFalse(lastUpgradeForCluster.isServiceCheckFailureAutoSkipped());
+
+    lastUpgradeForCluster.setAutoSkipComponentFailures(true);
+    lastUpgradeForCluster.setAutoSkipServiceCheckFailures(true);
+    dao.merge(lastUpgradeForCluster);
+
+    lastUpgradeForCluster = dao.findLastUpgradeForCluster(1);
+    Assert.assertTrue(lastUpgradeForCluster.isComponentFailureAutoSkipped());
+    Assert.assertTrue(lastUpgradeForCluster.isServiceCheckFailureAutoSkipped());
+  }
 }

ambari-server/src/test/java/org/apache/ambari/server/upgrade/UpgradeCatalog213Test.java

@@ -18,13 +18,26 @@
 
 package org.apache.ambari.server.upgrade;
 
-import com.google.inject.AbstractModule;
-import com.google.inject.Binder;
-import com.google.inject.Guice;
-import com.google.inject.Injector;
-import com.google.inject.Module;
-import com.google.inject.Provider;
-import com.google.inject.persist.PersistService;
+import static org.easymock.EasyMock.anyObject;
+import static org.easymock.EasyMock.capture;
+import static org.easymock.EasyMock.createMockBuilder;
+import static org.easymock.EasyMock.createNiceMock;
+import static org.easymock.EasyMock.createStrictMock;
+import static org.easymock.EasyMock.eq;
+import static org.easymock.EasyMock.expect;
+import static org.easymock.EasyMock.expectLastCall;
+import static org.easymock.EasyMock.replay;
+import static org.easymock.EasyMock.reset;
+import static org.easymock.EasyMock.verify;
+
+import java.lang.reflect.InvocationTargetException;
+import java.lang.reflect.Method;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+
+import javax.persistence.EntityManager;
+
 import org.apache.ambari.server.AmbariException;
 import org.apache.ambari.server.api.services.AmbariMetaInfo;
 import org.apache.ambari.server.configuration.Configuration;
@@ -51,24 +64,13 @@ import org.junit.Assert;
 import org.junit.Before;
 import org.junit.Test;
 
-import javax.persistence.EntityManager;
-import java.lang.reflect.InvocationTargetException;
-import java.lang.reflect.Method;
-import java.util.HashMap;
-import java.util.List;
-import java.util.Map;
-
-import static org.easymock.EasyMock.anyObject;
-import static org.easymock.EasyMock.capture;
-import static org.easymock.EasyMock.createMockBuilder;
-import static org.easymock.EasyMock.createNiceMock;
-import static org.easymock.EasyMock.createStrictMock;
-import static org.easymock.EasyMock.eq;
-import static org.easymock.EasyMock.expect;
-import static org.easymock.EasyMock.expectLastCall;
-import static org.easymock.EasyMock.replay;
-import static org.easymock.EasyMock.reset;
-import static org.easymock.EasyMock.verify;
+import com.google.inject.AbstractModule;
+import com.google.inject.Binder;
+import com.google.inject.Guice;
+import com.google.inject.Injector;
+import com.google.inject.Module;
+import com.google.inject.Provider;
+import com.google.inject.persist.PersistService;
 
 /**
  * {@link org.apache.ambari.server.upgrade.UpgradeCatalog213} unit tests.
@@ -108,6 +110,7 @@ public class UpgradeCatalog213Test {
     final DBAccessor dbAccessor = createNiceMock(DBAccessor.class);
 
     dbAccessor.addColumn(eq("upgrade"), anyObject(DBAccessor.DBColumnInfo.class));
+    expectLastCall().times(3);
 
     replay(dbAccessor);
     Module module = new Module() {
@@ -175,8 +178,8 @@ public class UpgradeCatalog213Test {
   }
 
   @Test
-  public void testPopulateDowngradeAllowed() throws Exception {
-    Method executeStackPreDMLUpdates = UpgradeCatalog213.class.getDeclaredMethod("populateDowngradeAllowed");
+  public void testExecuteUpgradePreDMLUpdates() throws Exception {
+    Method executeStackPreDMLUpdates = UpgradeCatalog213.class.getDeclaredMethod("executeUpgradePreDMLUpdates");
 
     final UpgradeCatalog213 upgradeCatalog213 = createMockBuilder(UpgradeCatalog213.class)
       .addMockedMethod(executeStackPreDMLUpdates).createMock();
@@ -190,7 +193,7 @@ public class UpgradeCatalog213Test {
       }
     });
 
-    upgradeCatalog213.populateDowngradeAllowed();
+    upgradeCatalog213.executeUpgradePreDMLUpdates();
     expectLastCall().once();
 
     replay(upgradeCatalog213);
@@ -587,6 +590,8 @@ public class UpgradeCatalog213Test {
 
     // executeUpgradeDDLUpdates
     mockedDbAccessor.addColumn(eq("upgrade"), capture(capturedColumn));
+    mockedDbAccessor.addColumn(eq("upgrade"), capture(capturedColumn));
+    mockedDbAccessor.addColumn(eq("upgrade"), capture(capturedColumn));
 
     // addKerberosDescriptorTable
     mockedDbAccessor.createTable(capture(capturedTableName), capture(capturedColumns), capture(capturedPKColumn));