AMBARI-8502 - Alerts: Upgrade Should Remove Nagios If Installed (jonathanhurley)

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

@@ -103,8 +103,6 @@ public class AlertResourceProvider extends ReadOnlyResourceProvider {
 
     // keys
     KEY_PROPERTY_IDS.put(Resource.Type.Cluster, ALERT_CLUSTER_NAME);
-    KEY_PROPERTY_IDS.put(Resource.Type.Service, ALERT_SERVICE);
-    KEY_PROPERTY_IDS.put(Resource.Type.Host, ALERT_HOST);
     KEY_PROPERTY_IDS.put(Resource.Type.Alert, ALERT_ID);
   }
 

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

@@ -20,11 +20,26 @@ package org.apache.ambari.server.upgrade;
 
 import java.sql.SQLException;
 import java.util.ArrayList;
+import java.util.Collection;
 import java.util.List;
 
 import org.apache.ambari.server.AmbariException;
 import org.apache.ambari.server.orm.DBAccessor;
 import org.apache.ambari.server.orm.DBAccessor.DBColumnInfo;
+import org.apache.ambari.server.orm.dao.ClusterDAO;
+import org.apache.ambari.server.orm.dao.ClusterServiceDAO;
+import org.apache.ambari.server.orm.dao.HostComponentDesiredStateDAO;
+import org.apache.ambari.server.orm.dao.HostComponentStateDAO;
+import org.apache.ambari.server.orm.dao.ServiceComponentDesiredStateDAO;
+import org.apache.ambari.server.orm.dao.ServiceDesiredStateDAO;
+import org.apache.ambari.server.orm.entities.ClusterEntity;
+import org.apache.ambari.server.orm.entities.ClusterServiceEntity;
+import org.apache.ambari.server.orm.entities.ClusterServiceEntityPK;
+import org.apache.ambari.server.orm.entities.HostComponentDesiredStateEntity;
+import org.apache.ambari.server.orm.entities.HostComponentStateEntity;
+import org.apache.ambari.server.orm.entities.ServiceComponentDesiredStateEntity;
+import org.apache.ambari.server.orm.entities.ServiceComponentDesiredStateEntityPK;
+import org.apache.ambari.server.orm.entities.ServiceDesiredStateEntity;
 import org.apache.ambari.server.state.SecurityState;
 import org.apache.ambari.server.state.UpgradeState;
 import org.slf4j.Logger;
@@ -174,8 +189,6 @@ public class UpgradeCatalog200 extends AbstractUpgradeCatalog {
     dbAccessor.executeQuery("INSERT INTO ambari_sequences(sequence_name, sequence_value) VALUES('cluster_version_id_seq', 0)", false);
     dbAccessor.executeQuery("INSERT INTO ambari_sequences(sequence_name, sequence_value) VALUES('host_version_id_seq', 0)", false);
 
-
-
     // upgrade tables
     columns = new ArrayList<DBColumnInfo>();
     columns.add(new DBAccessor.DBColumnInfo("upgrade_id", Long.class, null, null, false));
@@ -208,8 +221,6 @@ public class UpgradeCatalog200 extends AbstractUpgradeCatalog {
     dbAccessor.createTable("upgrade_item", columns, "upgrade_item_id");
     dbAccessor.addFKConstraint("upgrade_item", "fk_upgrade_item_upgrade_group_id", "upgrade_group_id", "upgrade_group", "upgrade_group_id", false);
     dbAccessor.executeQuery("INSERT INTO ambari_sequences(sequence_name, sequence_value) VALUES('upgrade_item_id_seq', 0)", false);
-
-
   }
 
   // ----- UpgradeCatalog ----------------------------------------------------
@@ -219,6 +230,84 @@ public class UpgradeCatalog200 extends AbstractUpgradeCatalog {
    */
   @Override
   protected void executeDMLUpdates() throws AmbariException, SQLException {
+    // remove NAGIOS to make way for the new embedded alert framework
+    removeNagiosService();
+  }
 
+  /**
+   * Removes Nagios and all associated components and states.
+   */
+  protected void removeNagiosService() {
+    executeInTransaction(new RemoveNagiosRunnable());
+  }
+
+  /**
+   * The RemoveNagiosRunnable is used to remove Nagios from the cluster. This
+   * runnable is exepected to run inside of a transation so that if any of the
+   * removals fails, Nagios is returned to a valid service state.
+   */
+  protected final class RemoveNagiosRunnable implements Runnable {
+
+    /**
+     * {@inheritDoc}
+     */
+    @Override
+    public void run() {
+      ClusterDAO clusterDao = injector.getInstance(ClusterDAO.class);
+      ClusterServiceDAO clusterServiceDao = injector.getInstance(ClusterServiceDAO.class);
+      ServiceComponentDesiredStateDAO componentDesiredStateDao = injector.getInstance(ServiceComponentDesiredStateDAO.class);
+      ServiceDesiredStateDAO desiredStateDao = injector.getInstance(ServiceDesiredStateDAO.class);
+      HostComponentDesiredStateDAO hostComponentDesiredStateDao = injector.getInstance(HostComponentDesiredStateDAO.class);
+      HostComponentStateDAO hostComponentStateDao = injector.getInstance(HostComponentStateDAO.class);
+
+      List<ClusterEntity> clusters = clusterDao.findAll();
+      if (null == clusters) {
+        return;
+      }
+
+      for (ClusterEntity cluster : clusters) {
+        ClusterServiceEntity nagios = clusterServiceDao.findByClusterAndServiceNames(
+            cluster.getClusterName(), "NAGIOS");
+
+        if (null == nagios) {
+          continue;
+        }
+
+        Collection<ServiceComponentDesiredStateEntity> serviceComponentDesiredStates = nagios.getServiceComponentDesiredStateEntities();
+        ServiceDesiredStateEntity serviceDesiredState = nagios.getServiceDesiredStateEntity();
+
+        // remove all component states
+        for (ServiceComponentDesiredStateEntity componentDesiredState : serviceComponentDesiredStates) {
+          Collection<HostComponentStateEntity> hostComponentStateEntities = componentDesiredState.getHostComponentStateEntities();
+          Collection<HostComponentDesiredStateEntity> hostComponentDesiredStateEntities = componentDesiredState.getHostComponentDesiredStateEntities();
+
+          // remove host states
+          for (HostComponentStateEntity hostComponentState : hostComponentStateEntities) {
+            hostComponentStateDao.remove(hostComponentState);
+          }
+
+          // remove host desired states
+          for (HostComponentDesiredStateEntity hostComponentDesiredState : hostComponentDesiredStateEntities) {
+            hostComponentDesiredStateDao.remove(hostComponentDesiredState);
+          }
+
+          // remove component state
+          ServiceComponentDesiredStateEntityPK primaryKey = new ServiceComponentDesiredStateEntityPK();
+          primaryKey.setClusterId(nagios.getClusterId());
+          primaryKey.setComponentName(componentDesiredState.getComponentName());
+          primaryKey.setServiceName(componentDesiredState.getServiceName());
+          componentDesiredStateDao.removeByPK(primaryKey);
+        }
+
+        // remove service state
+        desiredStateDao.remove(serviceDesiredState);
+
+        // remove service
+        ClusterServiceEntityPK primaryKey = new ClusterServiceEntityPK();
+        primaryKey.setClusterId(nagios.getClusterId());
+        primaryKey.setServiceName(nagios.getServiceName());
+        clusterServiceDao.removeByPK(primaryKey);
+      }
+    }
   }
 }

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

@@ -17,51 +17,46 @@
  */
 package org.apache.ambari.server.upgrade;
 
-import com.google.inject.Guice;
-import com.google.inject.Injector;
-import com.google.inject.persist.PersistService;
+import javax.persistence.EntityManager;
+
 import junit.framework.Assert;
+
 import org.apache.ambari.server.AmbariException;
 import org.apache.ambari.server.orm.GuiceJpaInitializer;
 import org.apache.ambari.server.orm.InMemoryDefaultTestModule;
 import org.apache.ambari.server.orm.dao.ClusterDAO;
-import org.apache.ambari.server.orm.dao.ClusterServiceDAO;
 import org.apache.ambari.server.orm.dao.HostComponentDesiredStateDAO;
-import org.apache.ambari.server.orm.dao.HostDAO;
 import org.apache.ambari.server.orm.dao.KeyValueDAO;
-import org.apache.ambari.server.orm.dao.ResourceTypeDAO;
 import org.apache.ambari.server.orm.entities.ClusterConfigEntity;
 import org.apache.ambari.server.orm.entities.ClusterConfigMappingEntity;
 import org.apache.ambari.server.orm.entities.ClusterEntity;
 import org.apache.ambari.server.orm.entities.ClusterServiceEntity;
 import org.apache.ambari.server.orm.entities.HostComponentDesiredStateEntity;
-import org.apache.ambari.server.orm.entities.HostComponentStateEntity;
 import org.apache.ambari.server.orm.entities.HostEntity;
 import org.apache.ambari.server.orm.entities.KeyValueEntity;
-import org.apache.ambari.server.orm.entities.ResourceEntity;
-import org.apache.ambari.server.orm.entities.ResourceTypeEntity;
 import org.apache.ambari.server.orm.entities.ServiceComponentDesiredStateEntity;
-import org.apache.ambari.server.orm.entities.ServiceDesiredStateEntity;
 import org.apache.ambari.server.state.HostComponentAdminState;
 import org.junit.After;
 import org.junit.Before;
 import org.junit.Test;
 
-import javax.persistence.EntityManager;
-import java.util.ArrayList;
-import java.util.Collections;
+import com.google.inject.Guice;
+import com.google.inject.Injector;
+import com.google.inject.persist.PersistService;
 
 public class UpgradeCatalog150Test {
   private Injector injector;
   private final String CLUSTER_NAME = "c1";
-  private final String SERVICE_NAME = "HDFS";
   private final String HOST_NAME = "h1";
   private final String DESIRED_STACK_VERSION = "{\"stackName\":\"HDP\",\"stackVersion\":\"1.3.4\"}";
 
+  private UpgradeCatalogHelper upgradeCatalogHelper;
+
   @Before
   public void setup() throws Exception {
     injector = Guice.createInjector(new InMemoryDefaultTestModule());
     injector.getInstance(GuiceJpaInitializer.class);
+    upgradeCatalogHelper = injector.getInstance(UpgradeCatalogHelper.class);
   }
 
   @After
@@ -69,101 +64,6 @@ public class UpgradeCatalog150Test {
     injector.getInstance(PersistService.class).stop();
   }
 
-  private ClusterEntity createCluster() {
-    ResourceTypeDAO resourceTypeDAO = injector.getInstance(ResourceTypeDAO.class);
-
-    // create an admin resource to represent this cluster
-    ResourceTypeEntity resourceTypeEntity = resourceTypeDAO.findById(ResourceTypeEntity.CLUSTER_RESOURCE_TYPE);
-    if (resourceTypeEntity == null) {
-      resourceTypeEntity = new ResourceTypeEntity();
-      resourceTypeEntity.setId(ResourceTypeEntity.CLUSTER_RESOURCE_TYPE);
-      resourceTypeEntity.setName(ResourceTypeEntity.CLUSTER_RESOURCE_TYPE_NAME);
-      resourceTypeEntity = resourceTypeDAO.merge(resourceTypeEntity);
-    }
-    ResourceEntity resourceEntity = new ResourceEntity();
-    resourceEntity.setResourceType(resourceTypeEntity);
-
-    ClusterDAO clusterDAO = injector.getInstance(ClusterDAO.class);
-    ClusterEntity clusterEntity = new ClusterEntity();
-    clusterEntity.setClusterId(1L);
-    clusterEntity.setClusterName(CLUSTER_NAME);
-    clusterEntity.setDesiredStackVersion(DESIRED_STACK_VERSION);
-    clusterEntity.setResource(resourceEntity);
-    clusterDAO.create(clusterEntity);
-    return clusterEntity;
-  }
-
-  private ClusterServiceEntity createService(ClusterEntity clusterEntity) {
-    ClusterServiceDAO clusterServiceDAO = injector.getInstance(ClusterServiceDAO.class);
-    ClusterServiceEntity clusterServiceEntity = new ClusterServiceEntity();
-    clusterServiceEntity.setClusterId(1L);
-    clusterServiceEntity.setClusterEntity(clusterEntity);
-    clusterServiceEntity.setServiceName(SERVICE_NAME);
-    clusterServiceDAO.create(clusterServiceEntity);
-    return clusterServiceEntity;
-  }
-  
-  private ClusterServiceEntity addService(ClusterEntity clusterEntity, String serviceName) {
-    ClusterDAO clusterDAO = injector.getInstance(ClusterDAO.class);
-    
-    ClusterServiceEntity clusterServiceEntity = new ClusterServiceEntity();
-    clusterServiceEntity.setClusterEntity(clusterEntity);
-    clusterServiceEntity.setServiceName(serviceName);
-
-    ServiceDesiredStateEntity serviceDesiredStateEntity = new ServiceDesiredStateEntity();
-    serviceDesiredStateEntity.setDesiredStackVersion(DESIRED_STACK_VERSION);
-    serviceDesiredStateEntity.setClusterServiceEntity(clusterServiceEntity);
-
-    clusterServiceEntity.setServiceDesiredStateEntity(serviceDesiredStateEntity);
-    clusterEntity.getClusterServiceEntities().add(clusterServiceEntity);
-
-    clusterDAO.merge(clusterEntity);
-    
-    return clusterServiceEntity;
-  }
-
-
-  private HostEntity createHost(ClusterEntity clusterEntity) {
-    HostDAO hostDAO = injector.getInstance(HostDAO.class);
-    ClusterDAO clusterDAO = injector.getInstance(ClusterDAO.class);
-    HostEntity hostEntity = new HostEntity();
-    hostEntity.setHostName(HOST_NAME);
-    hostEntity.setClusterEntities(Collections.singletonList(clusterEntity));
-    hostDAO.create(hostEntity);
-    clusterEntity.getHostEntities().add(hostEntity);
-    clusterDAO.merge(clusterEntity);
-    return hostEntity;
-  }
-  
-  private void addComponent(ClusterEntity clusterEntity, ClusterServiceEntity clusterServiceEntity, HostEntity hostEntity, String componentName) {
-    ServiceComponentDesiredStateEntity componentDesiredStateEntity = new ServiceComponentDesiredStateEntity();
-    componentDesiredStateEntity.setClusterServiceEntity(clusterServiceEntity);
-    componentDesiredStateEntity.setComponentName(componentName);
-    componentDesiredStateEntity.setHostComponentStateEntities(new ArrayList<HostComponentStateEntity>());
-
-    HostComponentDesiredStateEntity hostComponentDesiredStateEntity = new HostComponentDesiredStateEntity();
-    hostComponentDesiredStateEntity.setAdminState(HostComponentAdminState.INSERVICE);
-    hostComponentDesiredStateEntity.setServiceComponentDesiredStateEntity(componentDesiredStateEntity);
-    hostComponentDesiredStateEntity.setHostEntity(hostEntity);
-
-    HostComponentStateEntity hostComponentStateEntity = new HostComponentStateEntity();
-    hostComponentStateEntity.setHostEntity(hostEntity);
-    hostComponentStateEntity.setHostName(hostEntity.getHostName());
-    hostComponentStateEntity.setCurrentStackVersion(clusterEntity.getDesiredStackVersion());
-    hostComponentStateEntity.setServiceComponentDesiredStateEntity(componentDesiredStateEntity);
-
-    componentDesiredStateEntity.getHostComponentStateEntities().add(hostComponentStateEntity);
-    componentDesiredStateEntity.setHostComponentDesiredStateEntities(new ArrayList<HostComponentDesiredStateEntity>());
-    componentDesiredStateEntity.getHostComponentDesiredStateEntities().add(hostComponentDesiredStateEntity);
-
-    hostEntity.getHostComponentStateEntities().add(hostComponentStateEntity);
-    hostEntity.getHostComponentDesiredStateEntities().add(hostComponentDesiredStateEntity);
-    clusterServiceEntity.getServiceComponentDesiredStateEntities().add(componentDesiredStateEntity);
-    
-    ClusterServiceDAO clusterServiceDAO = injector.getInstance(ClusterServiceDAO.class);
-    clusterServiceDAO.merge(clusterServiceEntity);
-  }
-  
   protected void executeInTransaction(Runnable func) {
     EntityManager entityManager = injector.getProvider(EntityManager.class).get();
     if (entityManager.getTransaction().isActive()) { //already started, reuse
@@ -186,26 +86,34 @@ public class UpgradeCatalog150Test {
 
   @Test
   public void testAddHistoryServer() throws AmbariException {
-    final ClusterEntity clusterEntity = createCluster();
-    final ClusterServiceEntity clusterServiceEntityMR = addService(clusterEntity, "MAPREDUCE");
-    final HostEntity hostEntity = createHost(clusterEntity);
+    final ClusterEntity clusterEntity = upgradeCatalogHelper.createCluster(
+        injector, CLUSTER_NAME, DESIRED_STACK_VERSION);
+    final ClusterServiceEntity clusterServiceEntityMR = upgradeCatalogHelper.addService(
+        injector, clusterEntity, "MAPREDUCE", DESIRED_STACK_VERSION);
+    final HostEntity hostEntity = upgradeCatalogHelper.createHost(injector,
+        clusterEntity, HOST_NAME);
 
     executeInTransaction(new Runnable() {
       @Override
       public void run() {
-        addComponent(clusterEntity, clusterServiceEntityMR, hostEntity, "JOBTRACKER");
+        upgradeCatalogHelper.addComponent(injector, clusterEntity,
+            clusterServiceEntityMR, hostEntity, "JOBTRACKER",
+            DESIRED_STACK_VERSION);
       }
     });
-    
+
     UpgradeCatalog150 upgradeCatalog150 = injector.getInstance(UpgradeCatalog150.class);
     upgradeCatalog150.addHistoryServer();
   }
 
   @Test
   public void testProcessDecommissionedDatanodes() throws Exception {
-    ClusterEntity clusterEntity = createCluster();
-    ClusterServiceEntity clusterServiceEntity = createService(clusterEntity);
-    HostEntity hostEntity = createHost(clusterEntity);
+    ClusterEntity clusterEntity = upgradeCatalogHelper.createCluster(injector,
+        CLUSTER_NAME, DESIRED_STACK_VERSION);
+    ClusterServiceEntity clusterServiceEntity = upgradeCatalogHelper.createService(
+        injector, clusterEntity, "HDFS");
+    HostEntity hostEntity = upgradeCatalogHelper.createHost(injector,
+        clusterEntity, HOST_NAME);
 
     ServiceComponentDesiredStateEntity componentDesiredStateEntity =
       new ServiceComponentDesiredStateEntity();
@@ -271,8 +179,11 @@ public class UpgradeCatalog150Test {
   public void testAddMissingLog4jConfigs() throws Exception {
     ClusterDAO clusterDAO = injector.getInstance(ClusterDAO.class);
 
-    ClusterEntity clusterEntity = createCluster();
-    ClusterServiceEntity clusterServiceEntityMR = addService(clusterEntity, "HDFS");
+    ClusterEntity clusterEntity = upgradeCatalogHelper.createCluster(injector,
+        CLUSTER_NAME, DESIRED_STACK_VERSION);
+
+    ClusterServiceEntity clusterServiceEntityMR = upgradeCatalogHelper.addService(
+        injector, clusterEntity, "HDFS", DESIRED_STACK_VERSION);
 
     Long clusterId = clusterEntity.getClusterId();
 
@@ -302,10 +213,10 @@ public class UpgradeCatalog150Test {
     }
     Assert.assertFalse(failFlag);
   }
-  
+
   @Test
   public void testGetSourceVersion() {
     UpgradeCatalog150 upgradeCatalog150 = injector.getInstance(UpgradeCatalog150.class);
     Assert.assertNull(upgradeCatalog150.getSourceVersion());
-  }  
+  }
 }

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

@@ -18,13 +18,56 @@
 
 package org.apache.ambari.server.upgrade;
 
-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 com.google.inject.persist.Transactional;
+import static junit.framework.Assert.assertEquals;
+import static junit.framework.Assert.assertNotNull;
+import static junit.framework.Assert.assertNull;
+import static junit.framework.Assert.assertTrue;
+import static org.apache.ambari.server.upgrade.UpgradeCatalog170.CLUSTER_STATE_STACK_HDP_2_1;
+import static org.apache.ambari.server.upgrade.UpgradeCatalog170.JOBS_VIEW_NAME;
+import static org.apache.ambari.server.upgrade.UpgradeCatalog170.SHOW_JOBS_FOR_NON_ADMIN_KEY;
+import static org.apache.ambari.server.upgrade.UpgradeCatalog170.YARN_RESOURCEMANAGER_WEBAPP_ADDRESS_PROPERTY;
+import static org.apache.ambari.server.upgrade.UpgradeCatalog170.YARN_SITE;
+import static org.apache.ambari.server.upgrade.UpgradeCatalog170.YARN_TIMELINE_SERVICE_WEBAPP_ADDRESS_PROPERTY;
+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.getCurrentArguments;
+import static org.easymock.EasyMock.isA;
+import static org.easymock.EasyMock.replay;
+import static org.easymock.EasyMock.reset;
+import static org.easymock.EasyMock.verify;
+
+import java.lang.reflect.Field;
+import java.lang.reflect.Method;
+import java.sql.Connection;
+import java.sql.PreparedStatement;
+import java.sql.ResultSet;
+import java.sql.SQLException;
+import java.util.ArrayList;
+import java.util.Arrays;
+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 javax.persistence.EntityManager;
+import javax.persistence.EntityTransaction;
+import javax.persistence.TypedQuery;
+import javax.persistence.criteria.CriteriaBuilder;
+import javax.persistence.criteria.CriteriaQuery;
+import javax.persistence.criteria.Order;
+import javax.persistence.criteria.Path;
+import javax.persistence.criteria.Predicate;
+import javax.persistence.criteria.Root;
+import javax.persistence.metamodel.SingularAttribute;
+
 import org.apache.ambari.server.AmbariException;
 import org.apache.ambari.server.configuration.Configuration;
 import org.apache.ambari.server.controller.AmbariManagementController;
@@ -32,11 +75,9 @@ import org.apache.ambari.server.orm.DBAccessor;
 import org.apache.ambari.server.orm.GuiceJpaInitializer;
 import org.apache.ambari.server.orm.InMemoryDefaultTestModule;
 import org.apache.ambari.server.orm.dao.ClusterDAO;
-import org.apache.ambari.server.orm.dao.ClusterServiceDAO;
 import org.apache.ambari.server.orm.dao.ConfigGroupConfigMappingDAO;
 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.KeyValueDAO;
 import org.apache.ambari.server.orm.dao.PermissionDAO;
 import org.apache.ambari.server.orm.dao.PrincipalDAO;
@@ -65,10 +106,8 @@ import org.apache.ambari.server.orm.entities.PermissionEntity;
 import org.apache.ambari.server.orm.entities.PrincipalEntity;
 import org.apache.ambari.server.orm.entities.PrivilegeEntity;
 import org.apache.ambari.server.orm.entities.ResourceEntity;
-import org.apache.ambari.server.orm.entities.ResourceTypeEntity;
 import org.apache.ambari.server.orm.entities.ServiceComponentDesiredStateEntity;
 import org.apache.ambari.server.orm.entities.ServiceComponentDesiredStateEntityPK;
-import org.apache.ambari.server.orm.entities.ServiceDesiredStateEntity;
 import org.apache.ambari.server.orm.entities.UserEntity;
 import org.apache.ambari.server.orm.entities.ViewEntity;
 import org.apache.ambari.server.orm.entities.ViewInstanceEntity;
@@ -76,9 +115,7 @@ import org.apache.ambari.server.state.Cluster;
 import org.apache.ambari.server.state.Clusters;
 import org.apache.ambari.server.state.Config;
 import org.apache.ambari.server.state.ConfigHelper;
-import org.apache.ambari.server.state.HostComponentAdminState;
 import org.apache.ambari.server.state.StackId;
-import org.apache.ambari.server.state.State;
 import org.apache.ambari.server.view.ViewRegistry;
 import org.easymock.Capture;
 import org.easymock.IAnswer;
@@ -87,54 +124,12 @@ import org.junit.Assert;
 import org.junit.Before;
 import org.junit.Test;
 
-import javax.persistence.EntityManager;
-import javax.persistence.EntityTransaction;
-import javax.persistence.TypedQuery;
-import javax.persistence.criteria.CriteriaBuilder;
-import javax.persistence.criteria.CriteriaQuery;
-import javax.persistence.criteria.Order;
-import javax.persistence.criteria.Path;
-import javax.persistence.criteria.Predicate;
-import javax.persistence.criteria.Root;
-import javax.persistence.metamodel.SingularAttribute;
-import java.lang.reflect.Field;
-import java.lang.reflect.Method;
-import java.sql.Connection;
-import java.sql.PreparedStatement;
-import java.sql.ResultSet;
-import java.sql.SQLException;
-import java.util.ArrayList;
-import java.util.Arrays;
-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 static junit.framework.Assert.assertEquals;
-import static junit.framework.Assert.assertNotNull;
-import static junit.framework.Assert.assertNull;
-import static junit.framework.Assert.assertTrue;
-import static org.apache.ambari.server.upgrade.UpgradeCatalog170.CLUSTER_STATE_STACK_HDP_2_1;
-import static org.apache.ambari.server.upgrade.UpgradeCatalog170.JOBS_VIEW_NAME;
-import static org.apache.ambari.server.upgrade.UpgradeCatalog170.SHOW_JOBS_FOR_NON_ADMIN_KEY;
-import static org.apache.ambari.server.upgrade.UpgradeCatalog170.YARN_RESOURCEMANAGER_WEBAPP_ADDRESS_PROPERTY;
-import static org.apache.ambari.server.upgrade.UpgradeCatalog170.YARN_SITE;
-import static org.apache.ambari.server.upgrade.UpgradeCatalog170.YARN_TIMELINE_SERVICE_WEBAPP_ADDRESS_PROPERTY;
-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.getCurrentArguments;
-import static org.easymock.EasyMock.isA;
-import static org.easymock.EasyMock.replay;
-import static org.easymock.EasyMock.reset;
-import static org.easymock.EasyMock.verify;
+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;
 
 /**
  * UpgradeCatalog170 unit tests.
@@ -149,6 +144,7 @@ public class UpgradeCatalog170Test {
 
   Provider<EntityManager> entityManagerProvider = createStrictMock(Provider.class);
   EntityManager entityManager = createNiceMock(EntityManager.class);
+  private UpgradeCatalogHelper upgradeCatalogHelper;
 
   @Before
   public void init() {
@@ -157,6 +153,7 @@ public class UpgradeCatalog170Test {
     replay(entityManagerProvider);
     injector = Guice.createInjector(new InMemoryDefaultTestModule());
     injector.getInstance(GuiceJpaInitializer.class);
+    upgradeCatalogHelper = injector.getInstance(UpgradeCatalogHelper.class);
   }
 
   @After
@@ -164,123 +161,8 @@ public class UpgradeCatalog170Test {
     injector.getInstance(PersistService.class).stop();
   }
 
-  private ClusterEntity createCluster() {
-    ResourceTypeDAO resourceTypeDAO = injector.getInstance(ResourceTypeDAO.class);
-
-    // create an admin resource to represent this cluster
-    ResourceTypeEntity resourceTypeEntity = resourceTypeDAO.findById(ResourceTypeEntity.CLUSTER_RESOURCE_TYPE);
-    if (resourceTypeEntity == null) {
-      resourceTypeEntity = new ResourceTypeEntity();
-      resourceTypeEntity.setId(ResourceTypeEntity.CLUSTER_RESOURCE_TYPE);
-      resourceTypeEntity.setName(ResourceTypeEntity.CLUSTER_RESOURCE_TYPE_NAME);
-      resourceTypeEntity = resourceTypeDAO.merge(resourceTypeEntity);
-    }
-    ResourceEntity resourceEntity = new ResourceEntity();
-    resourceEntity.setResourceType(resourceTypeEntity);
-
-    ClusterDAO clusterDAO = injector.getInstance(ClusterDAO.class);
-    ClusterEntity clusterEntity = new ClusterEntity();
-    clusterEntity.setClusterId(1L);
-    clusterEntity.setClusterName(CLUSTER_NAME);
-    clusterEntity.setDesiredStackVersion(DESIRED_STACK_VERSION);
-    clusterEntity.setProvisioningState(State.INIT);
-    clusterEntity.setResource(resourceEntity);
-    clusterDAO.create(clusterEntity);
-    return clusterEntity;
-  }
-
-  private ClusterServiceEntity createService(ClusterEntity clusterEntity, String serviceName) {
-    ClusterServiceDAO clusterServiceDAO = injector.getInstance(ClusterServiceDAO.class);
-    ClusterServiceEntity clusterServiceEntity = new ClusterServiceEntity();
-    clusterServiceEntity.setClusterId(1L);
-    clusterServiceEntity.setClusterEntity(clusterEntity);
-    clusterServiceEntity.setServiceName(serviceName);
-    clusterServiceDAO.create(clusterServiceEntity);
-    return clusterServiceEntity;
-  }
-
-  private ClusterServiceEntity addService(ClusterEntity clusterEntity, String serviceName) {
-    ClusterDAO clusterDAO = injector.getInstance(ClusterDAO.class);
-
-    ClusterServiceEntity clusterServiceEntity = createService(clusterEntity, serviceName);
-
-    ServiceDesiredStateEntity serviceDesiredStateEntity = new ServiceDesiredStateEntity();
-    serviceDesiredStateEntity.setDesiredStackVersion(DESIRED_STACK_VERSION);
-    serviceDesiredStateEntity.setClusterId(1L);
-    serviceDesiredStateEntity.setServiceName(serviceName);
-    serviceDesiredStateEntity.setClusterServiceEntity(clusterServiceEntity);
-
-    clusterServiceEntity.setServiceDesiredStateEntity(serviceDesiredStateEntity);
-    clusterEntity.getClusterServiceEntities().add(clusterServiceEntity);
-
-    clusterDAO.merge(clusterEntity);
-
-    return clusterServiceEntity;
-  }
-
-
-  private HostEntity createHost(ClusterEntity clusterEntity) {
-    HostDAO hostDAO = injector.getInstance(HostDAO.class);
-    ClusterDAO clusterDAO = injector.getInstance(ClusterDAO.class);
-    HostEntity hostEntity = new HostEntity();
-    hostEntity.setHostName(HOST_NAME);
-    hostEntity.setClusterEntities(Collections.singletonList(clusterEntity));
-    hostDAO.create(hostEntity);
-    clusterEntity.getHostEntities().add(hostEntity);
-    clusterDAO.merge(clusterEntity);
-    return hostEntity;
-  }
-
-  @Transactional
-  private void addComponent(ClusterEntity clusterEntity, ClusterServiceEntity clusterServiceEntity, HostEntity hostEntity, String componentName) {
-    ServiceComponentDesiredStateEntity componentDesiredStateEntity = new ServiceComponentDesiredStateEntity();
-    componentDesiredStateEntity.setClusterServiceEntity(clusterServiceEntity);
-    componentDesiredStateEntity.setComponentName(componentName);
-    componentDesiredStateEntity.setServiceName(clusterServiceEntity.getServiceName());
-    componentDesiredStateEntity.setDesiredStackVersion(DESIRED_STACK_VERSION);
-    componentDesiredStateEntity.setClusterServiceEntity(clusterServiceEntity);
-    componentDesiredStateEntity.setClusterId(clusterServiceEntity.getClusterId());
-
-    HostComponentDesiredStateDAO hostComponentDesiredStateDAO = injector.getInstance(HostComponentDesiredStateDAO.class);
-    HostComponentDesiredStateEntity hostComponentDesiredStateEntity = new HostComponentDesiredStateEntity();
-    hostComponentDesiredStateEntity.setClusterId(clusterEntity.getClusterId());
-    hostComponentDesiredStateEntity.setHostName(hostEntity.getHostName());
-    hostComponentDesiredStateEntity.setComponentName(componentName);
-    hostComponentDesiredStateEntity.setServiceName(clusterServiceEntity.getServiceName());
-    hostComponentDesiredStateEntity.setAdminState(HostComponentAdminState.INSERVICE);
-    hostComponentDesiredStateEntity.setServiceComponentDesiredStateEntity(componentDesiredStateEntity);
-    hostComponentDesiredStateEntity.setHostEntity(hostEntity);
-    hostComponentDesiredStateDAO.create(hostComponentDesiredStateEntity);
-
-
-    HostComponentStateEntity hostComponentStateEntity = new HostComponentStateEntity();
-    hostComponentStateEntity.setHostEntity(hostEntity);
-    hostComponentStateEntity.setHostName(hostEntity.getHostName());
-    hostComponentStateEntity.setComponentName(componentName);
-    hostComponentStateEntity.setServiceName(clusterServiceEntity.getServiceName());
-    hostComponentStateEntity.setClusterId(clusterEntity.getClusterId());
-    hostComponentStateEntity.setCurrentStackVersion(clusterEntity.getDesiredStackVersion());
-    hostComponentStateEntity.setServiceComponentDesiredStateEntity(componentDesiredStateEntity);
-
-    componentDesiredStateEntity.setHostComponentStateEntities(Collections.singletonList(hostComponentStateEntity));
-    componentDesiredStateEntity.setHostComponentDesiredStateEntities(Collections.singletonList(hostComponentDesiredStateEntity));
-
-    hostEntity.getHostComponentStateEntities().add(hostComponentStateEntity);
-    hostEntity.getHostComponentDesiredStateEntities().add(hostComponentDesiredStateEntity);
-
-    clusterServiceEntity.getServiceComponentDesiredStateEntities().add(componentDesiredStateEntity);
-
-    ClusterServiceDAO clusterServiceDAO = injector.getInstance(ClusterServiceDAO.class);
-    ServiceComponentDesiredStateDAO serviceComponentDesiredStateDAO = injector.getInstance(ServiceComponentDesiredStateDAO.class);
-    HostDAO hostDAO = injector.getInstance(HostDAO.class);
-    serviceComponentDesiredStateDAO.merge(componentDesiredStateEntity);
-    hostDAO.merge(hostEntity);
-    clusterServiceDAO.merge(clusterServiceEntity);
-  }
-
   @Test
   public void testExecuteDDLUpdates() throws Exception {
-
     final DBAccessor dbAccessor = createNiceMock(DBAccessor.class);
     Connection connection = createNiceMock(Connection.class);
     PreparedStatement stmt = createNiceMock(PreparedStatement.class);
@@ -452,7 +334,7 @@ public class UpgradeCatalog170Test {
     Method n = AbstractUpgradeCatalog.class.getDeclaredMethod("getEntityManagerProvider");
     Method l = AbstractUpgradeCatalog.class.getDeclaredMethod
         ("addNewConfigurationsFromXml");
-    
+
     UpgradeCatalog170 upgradeCatalog = createMockBuilder(UpgradeCatalog170.class)
       .addMockedMethod(m).addMockedMethod(n).addMockedMethod(l).createMock();
 
@@ -655,16 +537,28 @@ public class UpgradeCatalog170Test {
 
   @Test
   public void testMoveHcatalogIntoHiveService()  throws AmbariException {
-    final ClusterEntity clusterEntity = createCluster();
-    final ClusterServiceEntity clusterServiceEntityHDFS = addService(clusterEntity, "HDFS");
-    final ClusterServiceEntity clusterServiceEntityHIVE = addService(clusterEntity, "HIVE");
-    final ClusterServiceEntity clusterServiceEntityHCATALOG = addService(clusterEntity, "HCATALOG");
-    final ClusterServiceEntity clusterServiceEntityWEBHCAT = addService(clusterEntity, "WEBHCAT");
-    final HostEntity hostEntity = createHost(clusterEntity);
-    addComponent(clusterEntity, clusterServiceEntityHDFS, hostEntity, "NAMENODE");
-    addComponent(clusterEntity, clusterServiceEntityHIVE, hostEntity, "HIVE_SERVER");
-    addComponent(clusterEntity, clusterServiceEntityHCATALOG, hostEntity, "HCAT");
-    addComponent(clusterEntity, clusterServiceEntityWEBHCAT, hostEntity, "WEBHCAT_SERVER");
+    final ClusterEntity clusterEntity = upgradeCatalogHelper.createCluster(
+        injector, CLUSTER_NAME, DESIRED_STACK_VERSION);
+    final ClusterServiceEntity clusterServiceEntityHDFS = upgradeCatalogHelper.addService(
+        injector, clusterEntity, "HDFS", DESIRED_STACK_VERSION);
+    final ClusterServiceEntity clusterServiceEntityHIVE = upgradeCatalogHelper.addService(
+        injector, clusterEntity, "HIVE", DESIRED_STACK_VERSION);
+    final ClusterServiceEntity clusterServiceEntityHCATALOG = upgradeCatalogHelper.addService(
+        injector, clusterEntity, "HCATALOG", DESIRED_STACK_VERSION);
+    final ClusterServiceEntity clusterServiceEntityWEBHCAT = upgradeCatalogHelper.addService(
+        injector, clusterEntity, "WEBHCAT", DESIRED_STACK_VERSION);
+    final HostEntity hostEntity = upgradeCatalogHelper.createHost(injector,
+        clusterEntity, HOST_NAME);
+    upgradeCatalogHelper.addComponent(injector, clusterEntity,
+        clusterServiceEntityHDFS, hostEntity, "NAMENODE", DESIRED_STACK_VERSION);
+    upgradeCatalogHelper.addComponent(injector, clusterEntity,
+        clusterServiceEntityHIVE, hostEntity, "HIVE_SERVER",
+        DESIRED_STACK_VERSION);
+    upgradeCatalogHelper.addComponent(injector, clusterEntity,
+        clusterServiceEntityHCATALOG, hostEntity, "HCAT", DESIRED_STACK_VERSION);
+    upgradeCatalogHelper.addComponent(injector, clusterEntity,
+        clusterServiceEntityWEBHCAT, hostEntity, "WEBHCAT_SERVER",
+        DESIRED_STACK_VERSION);
     UpgradeCatalog170 upgradeCatalog170 = injector.getInstance(UpgradeCatalog170.class);
     upgradeCatalog170.moveHcatalogIntoHiveService();
 
@@ -697,7 +591,8 @@ public class UpgradeCatalog170Test {
 
   @Test
   public void updateClusterProvisionState()  throws AmbariException {
-    ClusterEntity clusterEntity = createCluster();
+    ClusterEntity clusterEntity = upgradeCatalogHelper.createCluster(injector,
+        CLUSTER_NAME, DESIRED_STACK_VERSION);
     UpgradeCatalog170 upgradeCatalog170 = injector.getInstance(UpgradeCatalog170.class);
     upgradeCatalog170.updateClusterProvisionState();    //action
 

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

@@ -20,18 +20,22 @@ package org.apache.ambari.server.upgrade;
 
 import static junit.framework.Assert.assertEquals;
 import static junit.framework.Assert.assertFalse;
+import static junit.framework.Assert.assertNotNull;
 import static junit.framework.Assert.assertNull;
 import static junit.framework.Assert.assertTrue;
 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.Field;
+import java.lang.reflect.Method;
 import java.sql.ResultSet;
 import java.sql.SQLException;
 import java.util.List;
@@ -43,6 +47,19 @@ import org.apache.ambari.server.orm.DBAccessor;
 import org.apache.ambari.server.orm.DBAccessor.DBColumnInfo;
 import org.apache.ambari.server.orm.GuiceJpaInitializer;
 import org.apache.ambari.server.orm.InMemoryDefaultTestModule;
+import org.apache.ambari.server.orm.dao.ClusterServiceDAO;
+import org.apache.ambari.server.orm.dao.HostComponentDesiredStateDAO;
+import org.apache.ambari.server.orm.dao.HostComponentStateDAO;
+import org.apache.ambari.server.orm.dao.ServiceComponentDesiredStateDAO;
+import org.apache.ambari.server.orm.entities.ClusterEntity;
+import org.apache.ambari.server.orm.entities.ClusterServiceEntity;
+import org.apache.ambari.server.orm.entities.HostComponentDesiredStateEntity;
+import org.apache.ambari.server.orm.entities.HostComponentDesiredStateEntityPK;
+import org.apache.ambari.server.orm.entities.HostComponentStateEntity;
+import org.apache.ambari.server.orm.entities.HostComponentStateEntityPK;
+import org.apache.ambari.server.orm.entities.HostEntity;
+import org.apache.ambari.server.orm.entities.ServiceComponentDesiredStateEntity;
+import org.apache.ambari.server.orm.entities.ServiceComponentDesiredStateEntityPK;
 import org.apache.ambari.server.state.SecurityState;
 import org.easymock.Capture;
 import org.junit.After;
@@ -61,10 +78,14 @@ import com.google.inject.persist.PersistService;
  * {@link UpgradeCatalog200} unit tests.
  */
 public class UpgradeCatalog200Test {
+  private final String CLUSTER_NAME = "c1";
+  private final String HOST_NAME = "h1";
+  private final String DESIRED_STACK_VERSION = "{\"stackName\":\"HDP\",\"stackVersion\":\"2.0.6\"}";
 
   private Injector injector;
   private Provider<EntityManager> entityManagerProvider = createStrictMock(Provider.class);
   private EntityManager entityManager = createNiceMock(EntityManager.class);
+  private UpgradeCatalogHelper upgradeCatalogHelper;
 
   @Before
   public void init() {
@@ -73,6 +94,8 @@ public class UpgradeCatalog200Test {
     replay(entityManagerProvider);
     injector = Guice.createInjector(new InMemoryDefaultTestModule());
     injector.getInstance(GuiceJpaInitializer.class);
+
+    upgradeCatalogHelper = injector.getInstance(UpgradeCatalogHelper.class);
   }
 
   @After
@@ -194,8 +217,86 @@ public class UpgradeCatalog200Test {
     assertEquals(7, upgradeItemCapture.getValue().size());
   }
 
+  /**
+   * Tests that each DML method is invoked.
+   *
+   * @throws Exception
+   */
   @Test
   public void testExecuteDMLUpdates() throws Exception {
+    Method removeNagiosService = UpgradeCatalog200.class.getDeclaredMethod("removeNagiosService");
+
+    UpgradeCatalog200 upgradeCatalog = createMockBuilder(
+        UpgradeCatalog200.class).addMockedMethod(removeNagiosService).createMock();
+
+    upgradeCatalog.removeNagiosService();
+    expectLastCall().once();
+
+    replay(upgradeCatalog);
+
+    upgradeCatalog.executeDMLUpdates();
+
+    verify(upgradeCatalog);
+  }
+
+  /**
+   * Tests that Nagios is correctly removed.
+   *
+   * @throws Exception
+   */
+  @Test
+  public void testDeleteNagiosService() throws Exception {
+    final ClusterEntity clusterEntity = upgradeCatalogHelper.createCluster(
+        injector, CLUSTER_NAME, DESIRED_STACK_VERSION);
+
+    final ClusterServiceEntity clusterServiceEntityNagios = upgradeCatalogHelper.addService(
+        injector, clusterEntity, "NAGIOS", DESIRED_STACK_VERSION);
+
+    final HostEntity hostEntity = upgradeCatalogHelper.createHost(injector,
+        clusterEntity, HOST_NAME);
+
+    upgradeCatalogHelper.addComponent(injector, clusterEntity,
+        clusterServiceEntityNagios, hostEntity, "NAGIOS_SERVER",
+        DESIRED_STACK_VERSION);
+
+    UpgradeCatalog200 upgradeCatalog200 = injector.getInstance(UpgradeCatalog200.class);
+
+    ServiceComponentDesiredStateDAO serviceComponentDesiredStateDAO = injector.getInstance(ServiceComponentDesiredStateDAO.class);
+    ServiceComponentDesiredStateEntityPK pkNagiosServer = new ServiceComponentDesiredStateEntityPK();
+    pkNagiosServer.setComponentName("NAGIOS_SERVER");
+    pkNagiosServer.setClusterId(clusterEntity.getClusterId());
+    pkNagiosServer.setServiceName("NAGIOS");
+    ServiceComponentDesiredStateEntity serviceComponentDesiredStateEntity = serviceComponentDesiredStateDAO.findByPK(pkNagiosServer);
+    assertNotNull(serviceComponentDesiredStateEntity);
+
+    HostComponentDesiredStateDAO hostComponentDesiredStateDAO = injector.getInstance(HostComponentDesiredStateDAO.class);
+    HostComponentDesiredStateEntityPK hcDesiredStateEntityPk = new HostComponentDesiredStateEntityPK();
+    hcDesiredStateEntityPk.setServiceName("NAGIOS");
+    hcDesiredStateEntityPk.setClusterId(clusterEntity.getClusterId());
+    hcDesiredStateEntityPk.setComponentName("NAGIOS_SERVER");
+    hcDesiredStateEntityPk.setHostName(HOST_NAME);
+    HostComponentDesiredStateEntity hcDesiredStateEntity = hostComponentDesiredStateDAO.findByPK(hcDesiredStateEntityPk);
+    assertNotNull(hcDesiredStateEntity);
+
+    HostComponentStateDAO hostComponentStateDAO = injector.getInstance(HostComponentStateDAO.class);
+    HostComponentStateEntityPK hcStateEntityPk = new HostComponentStateEntityPK();
+    hcStateEntityPk.setServiceName("NAGIOS");
+    hcStateEntityPk.setClusterId(clusterEntity.getClusterId());
+    hcStateEntityPk.setComponentName("NAGIOS_SERVER");
+    hcStateEntityPk.setHostName(HOST_NAME);
+    HostComponentStateEntity hcStateEntity = hostComponentStateDAO.findByPK(hcStateEntityPk);
+    assertNotNull(hcStateEntity);
+
+    ClusterServiceDAO clusterServiceDao = injector.getInstance(ClusterServiceDAO.class);
+    ClusterServiceEntity clusterService = clusterServiceDao.findByClusterAndServiceNames(
+        CLUSTER_NAME, "NAGIOS");
+
+    upgradeCatalog200.removeNagiosService();
+
+    clusterService = clusterServiceDao.findByClusterAndServiceNames(
+        CLUSTER_NAME, "NAGIOS");
+
+    assertNull(clusterService);
   }
 
   /**

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

@@ -0,0 +1,215 @@
+/**
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.ambari.server.upgrade;
+
+import java.util.Collections;
+
+import org.apache.ambari.server.orm.dao.ClusterDAO;
+import org.apache.ambari.server.orm.dao.ClusterServiceDAO;
+import org.apache.ambari.server.orm.dao.HostComponentDesiredStateDAO;
+import org.apache.ambari.server.orm.dao.HostDAO;
+import org.apache.ambari.server.orm.dao.ResourceTypeDAO;
+import org.apache.ambari.server.orm.dao.ServiceComponentDesiredStateDAO;
+import org.apache.ambari.server.orm.entities.ClusterEntity;
+import org.apache.ambari.server.orm.entities.ClusterServiceEntity;
+import org.apache.ambari.server.orm.entities.HostComponentDesiredStateEntity;
+import org.apache.ambari.server.orm.entities.HostComponentStateEntity;
+import org.apache.ambari.server.orm.entities.HostEntity;
+import org.apache.ambari.server.orm.entities.ResourceEntity;
+import org.apache.ambari.server.orm.entities.ResourceTypeEntity;
+import org.apache.ambari.server.orm.entities.ServiceComponentDesiredStateEntity;
+import org.apache.ambari.server.orm.entities.ServiceDesiredStateEntity;
+import org.apache.ambari.server.state.HostComponentAdminState;
+import org.apache.ambari.server.state.State;
+
+import com.google.inject.Injector;
+import com.google.inject.persist.Transactional;
+
+/**
+ * The {@link UpgradeCatalogHelper} contains utility methods to help and of the
+ * {@link UpgradeCatalog} tests.
+ */
+public class UpgradeCatalogHelper {
+
+  /**
+   * Creates a cluster with the specified name and stack.
+   *
+   * @param injector
+   * @param clusterName
+   * @param desiredStackVersion
+   * @return
+   */
+  protected ClusterEntity createCluster(Injector injector, String clusterName,
+      String desiredStackVersion) {
+    ResourceTypeDAO resourceTypeDAO = injector.getInstance(ResourceTypeDAO.class);
+
+    // create an admin resource to represent this cluster
+    ResourceTypeEntity resourceTypeEntity = resourceTypeDAO.findById(ResourceTypeEntity.CLUSTER_RESOURCE_TYPE);
+    if (resourceTypeEntity == null) {
+      resourceTypeEntity = new ResourceTypeEntity();
+      resourceTypeEntity.setId(ResourceTypeEntity.CLUSTER_RESOURCE_TYPE);
+      resourceTypeEntity.setName(ResourceTypeEntity.CLUSTER_RESOURCE_TYPE_NAME);
+      resourceTypeEntity = resourceTypeDAO.merge(resourceTypeEntity);
+    }
+
+    ResourceEntity resourceEntity = new ResourceEntity();
+    resourceEntity.setResourceType(resourceTypeEntity);
+
+    ClusterDAO clusterDAO = injector.getInstance(ClusterDAO.class);
+    ClusterEntity clusterEntity = new ClusterEntity();
+    clusterEntity.setClusterId(1L);
+    clusterEntity.setClusterName(clusterName);
+    clusterEntity.setDesiredStackVersion(desiredStackVersion);
+    clusterEntity.setProvisioningState(State.INIT);
+    clusterEntity.setResource(resourceEntity);
+    clusterDAO.create(clusterEntity);
+    return clusterEntity;
+  }
+
+  /**
+   * Create a new service in the specified cluster.
+   *
+   * @param injector
+   * @param clusterEntity
+   * @param serviceName
+   * @return
+   */
+  protected ClusterServiceEntity createService(Injector injector,
+      ClusterEntity clusterEntity, String serviceName) {
+    ClusterServiceDAO clusterServiceDAO = injector.getInstance(ClusterServiceDAO.class);
+    ClusterServiceEntity clusterServiceEntity = new ClusterServiceEntity();
+    clusterServiceEntity.setClusterId(1L);
+    clusterServiceEntity.setClusterEntity(clusterEntity);
+    clusterServiceEntity.setServiceName(serviceName);
+    clusterServiceDAO.create(clusterServiceEntity);
+    return clusterServiceEntity;
+  }
+
+  /**
+   * Adds the specified service to a cluster. The service must already be part
+   * of the cluster.
+   *
+   * @param injector
+   * @param clusterEntity
+   * @param serviceName
+   * @param desiredStackVersion
+   * @return
+   */
+  protected ClusterServiceEntity addService(Injector injector,
+      ClusterEntity clusterEntity, String serviceName,
+      String desiredStackVersion) {
+    ClusterDAO clusterDAO = injector.getInstance(ClusterDAO.class);
+
+    ClusterServiceEntity clusterServiceEntity = createService(injector,
+        clusterEntity, serviceName);
+
+    ServiceDesiredStateEntity serviceDesiredStateEntity = new ServiceDesiredStateEntity();
+    serviceDesiredStateEntity.setDesiredStackVersion(desiredStackVersion);
+    serviceDesiredStateEntity.setClusterId(1L);
+    serviceDesiredStateEntity.setServiceName(serviceName);
+    serviceDesiredStateEntity.setClusterServiceEntity(clusterServiceEntity);
+
+    clusterServiceEntity.setServiceDesiredStateEntity(serviceDesiredStateEntity);
+    clusterEntity.getClusterServiceEntities().add(clusterServiceEntity);
+
+    clusterDAO.merge(clusterEntity);
+
+    return clusterServiceEntity;
+  }
+
+  /**
+   * Create a host in the specified cluster.
+   *
+   * @param injector
+   * @param clusterEntity
+   * @param hostName
+   * @return
+   */
+  protected HostEntity createHost(Injector injector,
+      ClusterEntity clusterEntity, String hostName) {
+    HostDAO hostDAO = injector.getInstance(HostDAO.class);
+    ClusterDAO clusterDAO = injector.getInstance(ClusterDAO.class);
+    HostEntity hostEntity = new HostEntity();
+    hostEntity.setHostName(hostName);
+    hostEntity.setClusterEntities(Collections.singletonList(clusterEntity));
+    hostDAO.create(hostEntity);
+    clusterEntity.getHostEntities().add(hostEntity);
+    clusterDAO.merge(clusterEntity);
+    return hostEntity;
+  }
+
+  /**
+   * Adds a host component for a given service and host.
+   *
+   * @param injector
+   * @param clusterEntity
+   * @param clusterServiceEntity
+   * @param hostEntity
+   * @param componentName
+   * @param desiredStackVersion
+   */
+  @Transactional
+  protected void addComponent(Injector injector, ClusterEntity clusterEntity,
+      ClusterServiceEntity clusterServiceEntity, HostEntity hostEntity,
+      String componentName, String desiredStackVersion) {
+    ServiceComponentDesiredStateEntity componentDesiredStateEntity = new ServiceComponentDesiredStateEntity();
+    componentDesiredStateEntity.setClusterServiceEntity(clusterServiceEntity);
+    componentDesiredStateEntity.setComponentName(componentName);
+    componentDesiredStateEntity.setServiceName(clusterServiceEntity.getServiceName());
+    componentDesiredStateEntity.setDesiredStackVersion(desiredStackVersion);
+    componentDesiredStateEntity.setClusterServiceEntity(clusterServiceEntity);
+    componentDesiredStateEntity.setClusterId(clusterServiceEntity.getClusterId());
+
+    HostComponentDesiredStateDAO hostComponentDesiredStateDAO = injector.getInstance(HostComponentDesiredStateDAO.class);
+    HostComponentDesiredStateEntity hostComponentDesiredStateEntity = new HostComponentDesiredStateEntity();
+    hostComponentDesiredStateEntity.setClusterId(clusterEntity.getClusterId());
+    hostComponentDesiredStateEntity.setHostName(hostEntity.getHostName());
+    hostComponentDesiredStateEntity.setComponentName(componentName);
+    hostComponentDesiredStateEntity.setServiceName(clusterServiceEntity.getServiceName());
+    hostComponentDesiredStateEntity.setAdminState(HostComponentAdminState.INSERVICE);
+    hostComponentDesiredStateEntity.setServiceComponentDesiredStateEntity(componentDesiredStateEntity);
+    hostComponentDesiredStateEntity.setHostEntity(hostEntity);
+    hostComponentDesiredStateDAO.create(hostComponentDesiredStateEntity);
+
+    HostComponentStateEntity hostComponentStateEntity = new HostComponentStateEntity();
+    hostComponentStateEntity.setHostEntity(hostEntity);
+    hostComponentStateEntity.setHostName(hostEntity.getHostName());
+    hostComponentStateEntity.setComponentName(componentName);
+    hostComponentStateEntity.setServiceName(clusterServiceEntity.getServiceName());
+    hostComponentStateEntity.setClusterId(clusterEntity.getClusterId());
+    hostComponentStateEntity.setCurrentStackVersion(clusterEntity.getDesiredStackVersion());
+    hostComponentStateEntity.setServiceComponentDesiredStateEntity(componentDesiredStateEntity);
+
+    componentDesiredStateEntity.setHostComponentStateEntities(Collections.singletonList(hostComponentStateEntity));
+    componentDesiredStateEntity.setHostComponentDesiredStateEntities(Collections.singletonList(hostComponentDesiredStateEntity));
+
+    hostEntity.getHostComponentStateEntities().add(hostComponentStateEntity);
+    hostEntity.getHostComponentDesiredStateEntities().add(
+        hostComponentDesiredStateEntity);
+
+    clusterServiceEntity.getServiceComponentDesiredStateEntities().add(
+        componentDesiredStateEntity);
+
+    ClusterServiceDAO clusterServiceDAO = injector.getInstance(ClusterServiceDAO.class);
+    ServiceComponentDesiredStateDAO serviceComponentDesiredStateDAO = injector.getInstance(ServiceComponentDesiredStateDAO.class);
+    HostDAO hostDAO = injector.getInstance(HostDAO.class);
+    serviceComponentDesiredStateDAO.merge(componentDesiredStateEntity);
+    hostDAO.merge(hostEntity);
+    clusterServiceDAO.merge(clusterServiceEntity);
+  }
+}