/*
 * 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.serveraction.upgrades;

import static org.junit.Assert.assertEquals;
import static org.junit.Assert.assertNotNull;
import static org.junit.Assert.assertTrue;

import java.io.File;
import java.io.FileInputStream;
import java.lang.reflect.Field;
import java.nio.charset.Charset;
import java.util.Collection;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

import jakarta.persistence.EntityManager;

import org.apache.ambari.server.AmbariException;
import org.apache.ambari.server.H2DatabaseCleaner;
import org.apache.ambari.server.ServiceComponentNotFoundException;
import org.apache.ambari.server.ServiceNotFoundException;
import org.apache.ambari.server.actionmanager.ExecutionCommandWrapper;
import org.apache.ambari.server.actionmanager.HostRoleCommand;
import org.apache.ambari.server.actionmanager.HostRoleCommandFactory;
import org.apache.ambari.server.actionmanager.HostRoleStatus;
import org.apache.ambari.server.agent.CommandReport;
import org.apache.ambari.server.agent.ExecutionCommand;
import org.apache.ambari.server.api.services.AmbariMetaInfo;
import org.apache.ambari.server.controller.AmbariManagementController;
import org.apache.ambari.server.controller.AmbariServer;
import org.apache.ambari.server.controller.ServiceConfigVersionResponse;
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.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;
import org.apache.ambari.server.orm.dao.RequestDAO;
import org.apache.ambari.server.orm.dao.StackDAO;
import org.apache.ambari.server.orm.dao.UpgradeDAO;
import org.apache.ambari.server.orm.entities.HostComponentStateEntity;
import org.apache.ambari.server.orm.entities.HostVersionEntity;
import org.apache.ambari.server.orm.entities.RepositoryVersionEntity;
import org.apache.ambari.server.orm.entities.RequestEntity;
import org.apache.ambari.server.orm.entities.StackEntity;
import org.apache.ambari.server.orm.entities.UpgradeEntity;
import org.apache.ambari.server.orm.entities.UpgradeHistoryEntity;
import org.apache.ambari.server.serveraction.ServerAction;
import org.apache.ambari.server.stack.upgrade.Direction;
import org.apache.ambari.server.stack.upgrade.UpgradePack;
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.ConfigFactory;
import org.apache.ambari.server.state.Host;
import org.apache.ambari.server.state.RepositoryVersionState;
import org.apache.ambari.server.state.Service;
import org.apache.ambari.server.state.ServiceComponent;
import org.apache.ambari.server.state.ServiceComponentFactory;
import org.apache.ambari.server.state.ServiceComponentHost;
import org.apache.ambari.server.state.ServiceComponentHostFactory;
import org.apache.ambari.server.state.ServiceFactory;
import org.apache.ambari.server.state.StackId;
import org.apache.ambari.server.state.State;
import org.apache.ambari.server.state.UpgradeState;
import org.apache.ambari.server.utils.EventBusSynchronizer;
import org.apache.ambari.spi.RepositoryType;
import org.apache.ambari.spi.upgrade.UpgradeType;
import org.apache.commons.io.IOUtils;
import org.apache.commons.lang.StringUtils;
import org.junit.After;
import org.junit.Assert;
import org.junit.Before;
import org.junit.Test;

import com.google.inject.Guice;
import com.google.inject.Inject;
import com.google.inject.Injector;
import com.google.inject.persist.UnitOfWork;

/**
 * Tests upgrade-related server side actions
 */
public class UpgradeActionTest {
  private static final String clusterName = "c1";

  private static final String HDP_2_1_1_0 = "2.1.1.0-1";
  private static final String HDP_2_1_1_1 = "2.1.1.1-2";

  private static final String HDP_2_2_0_1 = "2.2.0.1-3";

  private static final StackId HDP_21_STACK = new StackId("HDP-2.1.1");
  private static final StackId HDP_22_STACK = new StackId("HDP-2.2.0");

  private RepositoryVersionEntity sourceRepositoryVersion;

  private Injector m_injector;

  private AmbariManagementController amc;
  @Inject
  private OrmTestHelper m_helper;
  @Inject
  private RepositoryVersionDAO repoVersionDAO;
  @Inject
  private Clusters clusters;
  @Inject
  private HostVersionDAO hostVersionDAO;
  @Inject
  private HostDAO hostDAO;
  @Inject
  private HostRoleCommandFactory hostRoleCommandFactory;
  @Inject
  private ServiceFactory serviceFactory;
  @Inject
  private ServiceComponentFactory serviceComponentFactory;
  @Inject
  private ServiceComponentHostFactory serviceComponentHostFactory;
  @Inject
  private RequestDAO requestDAO;
  @Inject
  private UpgradeDAO upgradeDAO;
  @Inject
  private StackDAO stackDAO;
  @Inject
  private AmbariMetaInfo ambariMetaInfo;
  @Inject
  private FinalizeUpgradeAction finalizeUpgradeAction;
  @Inject
  private ConfigFactory configFactory;

  @Inject
  private RepositoryVersionDAO repositoryVersionDAO;

  @Inject
  private HostComponentStateDAO hostComponentStateDAO;

  private RepositoryVersionEntity repositoryVersion2110;
  private RepositoryVersionEntity repositoryVersion2111;
  private RepositoryVersionEntity repositoryVersion2201;

  @Before
  public void setup() throws Exception {
    m_injector = Guice.createInjector(new InMemoryDefaultTestModule());
    EventBusSynchronizer.synchronizeAmbariEventPublisher(m_injector);
    m_injector.getInstance(GuiceJpaInitializer.class);
    m_injector.injectMembers(this);
    m_injector.getInstance(UnitOfWork.class).begin();

    // Initialize AmbariManagementController
    amc = m_injector.getInstance(AmbariManagementController.class);

    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
  public void teardown() throws Exception {
    m_injector.getInstance(UnitOfWork.class).end();
    H2DatabaseCleaner.clearDatabase(m_injector.getProvider(EntityManager.class).get());
  }

  private void makeDowngradeCluster(RepositoryVersionEntity sourceRepoVersion,
      RepositoryVersionEntity targetRepoVersion) throws Exception {
    String hostName = "h1";

    clusters.addCluster(clusterName, sourceRepoVersion.getStackId());

    // add a host component
    clusters.addHost(hostName);

    Host host = clusters.getHost(hostName);

    Map<String, String> hostAttributes = new HashMap<>();
    hostAttributes.put("os_family", "redhat");
    hostAttributes.put("os_release_version", "6");
    host.setHostAttributes(hostAttributes);

    HostVersionEntity entity = new HostVersionEntity();
    entity.setHostEntity(hostDAO.findByName(hostName));
    entity.setRepositoryVersion(targetRepoVersion);
    entity.setState(RepositoryVersionState.INSTALLING);
    hostVersionDAO.create(entity);
  }

  private Cluster createUpgradeCluster(
      RepositoryVersionEntity sourceRepoVersion, String hostName) throws Exception {

    clusters.addCluster(clusterName, sourceRepoVersion.getStackId());
    Cluster cluster = clusters.getCluster(clusterName);

    // add a host component
    clusters.addHost(hostName);

    Host host = clusters.getHost(hostName);

    Map<String, String> hostAttributes = new HashMap<>();
    hostAttributes.put("os_family", "redhat");
    hostAttributes.put("os_release_version", "6");
    host.setHostAttributes(hostAttributes);

    // without this, HostEntity will not have a relation to ClusterEntity
    clusters.mapHostToCluster(hostName, clusterName);

    HostVersionEntity entity = new HostVersionEntity(hostDAO.findByName(hostName),
        sourceRepoVersion, RepositoryVersionState.INSTALLED);

    hostVersionDAO.create(entity);

    return cluster;
  }

  private void createHostVersions(RepositoryVersionEntity targetRepoVersion,
      String hostName) throws AmbariException {
    Cluster c = clusters.getCluster(clusterName);

    // create a single host with the UPGRADED HostVersionEntity
    HostDAO hostDAO = m_injector.getInstance(HostDAO.class);

    HostVersionEntity entity = new HostVersionEntity(hostDAO.findByName(hostName),
        targetRepoVersion, RepositoryVersionState.INSTALLED);

    hostVersionDAO.create(entity);

    // verify the UPGRADED host versions were created successfully
    List<HostVersionEntity> hostVersions = hostVersionDAO.findHostVersionByClusterAndRepository(
        c.getClusterId(), targetRepoVersion);

    assertEquals(1, hostVersions.size());
    assertEquals(RepositoryVersionState.INSTALLED, hostVersions.get(0).getState());
  }

  private void makeCrossStackUpgradeClusterAndSourceRepo(StackId sourceStack, String sourceRepo,
                                                         String hostName)throws Exception {

    clusters.addCluster(clusterName, sourceStack);

    StackEntity stackEntitySource = stackDAO.find(sourceStack.getStackName(), sourceStack.getStackVersion());

    assertNotNull(stackEntitySource);

    Cluster c = clusters.getCluster(clusterName);
    c.setCurrentStackVersion(sourceStack);
    c.setDesiredStackVersion(sourceStack);

    // add a host component
    clusters.addHost(hostName);
    Host host = clusters.getHost(hostName);

    Map<String, String> hostAttributes = new HashMap<>();
    hostAttributes.put("os_family", "redhat");
    hostAttributes.put("os_release_version", "6");
    host.setHostAttributes(hostAttributes);

    clusters.mapHostToCluster(hostName, clusterName);

    clusters.updateHostMappings(host);

    // Create the starting repo version
    sourceRepositoryVersion = m_helper.getOrCreateRepositoryVersion(sourceStack, sourceRepo);
  }

  private void makeCrossStackUpgradeTargetRepo(StackId targetStack, String targetRepo, String hostName) throws Exception{
    StackEntity stackEntityTarget = stackDAO.find(targetStack.getStackName(), targetStack.getStackVersion());
    assertNotNull(stackEntityTarget);

    m_helper.getOrCreateRepositoryVersion(new StackId(stackEntityTarget), targetRepo);

    // Start upgrading the newer repo

    HostDAO hostDAO = m_injector.getInstance(HostDAO.class);

    HostVersionEntity entity = new HostVersionEntity();
    entity.setHostEntity(hostDAO.findByName(hostName));
    entity.setRepositoryVersion(repoVersionDAO.findByStackAndVersion(targetStack, targetRepo));
    entity.setState(RepositoryVersionState.INSTALLED);
    hostVersionDAO.create(entity);
  }

  /***
   * During an Express Upgrade that crosses a stack version, Ambari calls UpdateDesiredRepositoryAction
   * in order to change the stack and apply configs.
   * The configs that are applied must be saved with the username that is passed in the role params.
   */
  @Test
  public void testExpressUpgradeUpdateDesiredRepositoryAction() throws Exception {
    StackId sourceStack = HDP_21_STACK;
    StackId targetStack = HDP_22_STACK;
    String sourceRepo = HDP_2_1_1_0;
    String hostName = "h1";

    // Must be a NON_ROLLING upgrade that jumps stacks in order for it to apply config changes.
    // That upgrade pack has changes for ZK and NameNode.
    String upgradePackName = "upgrade_nonrolling_new_stack";

    Map<String, UpgradePack> packs = ambariMetaInfo.getUpgradePacks(sourceStack.getStackName(), sourceStack.getStackVersion());
    Assert.assertTrue(packs.containsKey(upgradePackName));

    makeCrossStackUpgradeClusterAndSourceRepo(sourceStack, sourceRepo, hostName);

    Cluster cluster = clusters.getCluster(clusterName);

    // Install ZK and HDFS with some components
    Service zk = installService(cluster, "ZOOKEEPER", repositoryVersion2110);
    addServiceComponent(cluster, zk, "ZOOKEEPER_SERVER");
    addServiceComponent(cluster, zk, "ZOOKEEPER_CLIENT");
    createNewServiceComponentHost(cluster, "ZOOKEEPER", "ZOOKEEPER_SERVER", "h1");
    createNewServiceComponentHost(cluster, "ZOOKEEPER", "ZOOKEEPER_CLIENT", "h1");

    Service hdfs = installService(cluster, "HDFS", repositoryVersion2110);
    addServiceComponent(cluster, hdfs, "NAMENODE");
    addServiceComponent(cluster, hdfs, "DATANODE");
    createNewServiceComponentHost(cluster, "HDFS", "NAMENODE", "h1");
    createNewServiceComponentHost(cluster, "HDFS", "DATANODE", "h1");

    makeCrossStackUpgradeTargetRepo(targetStack, repositoryVersion2201.getVersion(), hostName);
    createUpgrade(cluster, repositoryVersion2201);

    Assert.assertNotNull(repositoryVersion2201);

    // Create some configs
    createConfigs(cluster);
    Collection<Config> configs = cluster.getAllConfigs();
    Assert.assertFalse(configs.isEmpty());

    Map<String, String> commandParams = new HashMap<>();
    ExecutionCommand executionCommand = new ExecutionCommand();
    executionCommand.setCommandParams(commandParams);
    Map<String, String> roleParams = new HashMap<>();

    // User that is performing the config changes
    String userName = "admin";
    roleParams.put(ServerAction.ACTION_USER_NAME, userName);
    executionCommand.setRoleParams(roleParams);
    executionCommand.setClusterName(clusterName);

    HostRoleCommand hostRoleCommand = hostRoleCommandFactory.create(null, null, null, null);
    hostRoleCommand.setExecutionCommandWrapper(new ExecutionCommandWrapper(executionCommand));

    // Call the action to change the desired stack and apply the configs from the Config Pack called by the Upgrade Pack.
    UpdateDesiredRepositoryAction action = m_injector.getInstance(UpdateDesiredRepositoryAction.class);
    action.setExecutionCommand(executionCommand);
    action.setHostRoleCommand(hostRoleCommand);

    List<ServiceConfigVersionResponse> configVersionsBefore = cluster.getServiceConfigVersions();

    CommandReport report = action.execute(null);
    assertNotNull(report);

    assertEquals(HostRoleStatus.COMPLETED.name(), report.getStatus());

    List<ServiceConfigVersionResponse> configVersionsAfter = cluster.getServiceConfigVersions();
    Assert.assertFalse(configVersionsAfter.isEmpty());

    assertTrue(configVersionsAfter.size() - configVersionsBefore.size() >= 1);
  }

  @Test
  public void testFinalizeDowngrade() throws Exception {
    makeDowngradeCluster(repositoryVersion2110, repositoryVersion2111);

    Cluster cluster = clusters.getCluster(clusterName);

    createUpgrade(cluster, repositoryVersion2111);

    Map<String, String> commandParams = new HashMap<>();
    ExecutionCommand executionCommand = new ExecutionCommand();
    executionCommand.setCommandParams(commandParams);
    executionCommand.setClusterName(clusterName);

    HostRoleCommand hostRoleCommand = hostRoleCommandFactory.create(null, null, null, null);
    hostRoleCommand.setExecutionCommandWrapper(new ExecutionCommandWrapper(executionCommand));

    finalizeUpgradeAction.setExecutionCommand(executionCommand);
    finalizeUpgradeAction.setHostRoleCommand(hostRoleCommand);

    CommandReport report = finalizeUpgradeAction.execute(null);
    assertNotNull(report);
    assertEquals(HostRoleStatus.COMPLETED.name(), report.getStatus());

    for (HostVersionEntity entity : hostVersionDAO.findByClusterAndHost(clusterName, "h1")) {
      if (StringUtils.equals(entity.getRepositoryVersion().getVersion(), repositoryVersion2110.getVersion())) {
        assertEquals(RepositoryVersionState.CURRENT, entity.getState());
      } else {
        assertEquals(RepositoryVersionState.INSTALLED, entity.getState());
      }
    }
  }

  @Test
  public void testFinalizeUpgrade() throws Exception {
    String hostName = "h1";

    createUpgradeCluster(repositoryVersion2110, hostName);
    createHostVersions(repositoryVersion2111, hostName);

    Cluster cluster = clusters.getCluster(clusterName);

    createUpgrade(cluster, repositoryVersion2111);

    // Finalize the upgrade
    Map<String, String> commandParams = new HashMap<>();
    ExecutionCommand executionCommand = new ExecutionCommand();
    executionCommand.setCommandParams(commandParams);
    executionCommand.setClusterName(clusterName);

    HostRoleCommand hostRoleCommand = hostRoleCommandFactory.create(null, null, null, null);
    hostRoleCommand.setExecutionCommandWrapper(new ExecutionCommandWrapper(executionCommand));

    finalizeUpgradeAction.setExecutionCommand(executionCommand);
    finalizeUpgradeAction.setHostRoleCommand(hostRoleCommand);

    // this should fail since the host versions have not moved to current
    CommandReport report = finalizeUpgradeAction.execute(null);
    assertEquals(HostRoleStatus.FAILED.name(), report.getStatus());

    List<HostVersionEntity> hostVersions = hostVersionDAO.findHostVersionByClusterAndRepository(
        cluster.getClusterId(), repositoryVersion2111);

    for (HostVersionEntity hostVersion : hostVersions) {
      hostVersion.setState(RepositoryVersionState.CURRENT);
    }

    report = finalizeUpgradeAction.execute(null);
    assertEquals(HostRoleStatus.COMPLETED.name(), report.getStatus());

    hostVersions = hostVersionDAO.findHostVersionByClusterAndRepository(cluster.getClusterId(),
        repositoryVersion2111);

    for (HostVersionEntity hostVersion : hostVersions) {
      Collection<HostComponentStateEntity> hostComponentStates = hostComponentStateDAO.findByHost(hostVersion.getHostName());
      for (HostComponentStateEntity hostComponentStateEntity: hostComponentStates) {
       assertEquals(UpgradeState.NONE, hostComponentStateEntity.getUpgradeState());
      }
    }
  }

  /**
   * Tests that finalize still works when there are hosts which are already
   * {@link RepositoryVersionState#CURRENT}.
   *
   * @throws Exception
   */
  @Test
  public void testFinalizeWithHostsAlreadyCurrent() throws Exception {
    String hostName = "h1";

    createUpgradeCluster(repositoryVersion2110, hostName);
    createHostVersions(repositoryVersion2111, hostName);

    // move the old version from CURRENT to INSTALLED and the new version from
    // UPGRADED to CURRENT - this will simulate what happens when a host is
    // removed before finalization and all hosts transition to CURRENT
    List<HostVersionEntity> hostVersions = hostVersionDAO.findAll();
    for (HostVersionEntity hostVersion : hostVersions) {
      if (hostVersion.getState() == RepositoryVersionState.CURRENT) {
        hostVersion.setState(RepositoryVersionState.INSTALLED);
      } else {
        hostVersion.setState(RepositoryVersionState.CURRENT);
      }

      hostVersionDAO.merge(hostVersion);
    }

    // Verify the repo before calling Finalize
    Cluster cluster = clusters.getCluster(clusterName);

    createUpgrade(cluster, repositoryVersion2111);

    // Finalize the upgrade
    Map<String, String> commandParams = new HashMap<>();

    ExecutionCommand executionCommand = new ExecutionCommand();
    executionCommand.setCommandParams(commandParams);
    executionCommand.setClusterName(clusterName);

    HostRoleCommand hostRoleCommand = hostRoleCommandFactory.create(null, null, null, null);
    hostRoleCommand.setExecutionCommandWrapper(new ExecutionCommandWrapper(executionCommand));

    finalizeUpgradeAction.setExecutionCommand(executionCommand);
    finalizeUpgradeAction.setHostRoleCommand(hostRoleCommand);

    CommandReport report = finalizeUpgradeAction.execute(null);
    assertNotNull(report);
    assertEquals(HostRoleStatus.COMPLETED.name(), report.getStatus());
  }

  /**
   * Tests that all host versions are correct after upgrade. This test will
   * ensure that the prior CURRENT versions are moved to INSTALLED while not
   * touching any others.
   *
   * @throws Exception
   */
  @Test
  public void testHostVersionsAfterUpgrade() throws Exception {
    String hostName = "h1";
    Cluster cluster = createUpgradeCluster(repositoryVersion2110, hostName);
    createHostVersions(repositoryVersion2111, hostName);
    createHostVersions(repositoryVersion2201, hostName);

    // Install ZK with some components
    Service zk = installService(cluster, "ZOOKEEPER", repositoryVersion2110);
    addServiceComponent(cluster, zk, "ZOOKEEPER_SERVER");
    addServiceComponent(cluster, zk, "ZOOKEEPER_CLIENT");
    createNewServiceComponentHost(cluster, "ZOOKEEPER", "ZOOKEEPER_SERVER", hostName);
    createNewServiceComponentHost(cluster, "ZOOKEEPER", "ZOOKEEPER_CLIENT", hostName);

    List<HostVersionEntity> hostVersions = hostVersionDAO.findAll();
    assertEquals(3, hostVersions.size());

    // repo 2110 - CURRENT (upgrading from)
    // repo 2111 - CURRENT (all hosts reported in during upgrade)
    // repo 2201 - NOT_REQUIRED (different stack)
    for (HostVersionEntity hostVersion : hostVersions) {
      RepositoryVersionEntity hostRepoVersion = hostVersion.getRepositoryVersion();
      if (repositoryVersion2110.equals(hostRepoVersion)) {
        hostVersion.setState(RepositoryVersionState.CURRENT);
      } else if (repositoryVersion2111.equals(hostRepoVersion)) {
        hostVersion.setState(RepositoryVersionState.CURRENT);
      } else {
        hostVersion.setState(RepositoryVersionState.NOT_REQUIRED);
      }

      hostVersionDAO.merge(hostVersion);
    }

    // upgrade to 2111
    createUpgrade(cluster, repositoryVersion2111);

    // push all services to the correct repo version for finalize
    Map<String, Service> services = cluster.getServices();
    assertTrue(services.size() > 0);
    for (Service service : services.values()) {
      service.setDesiredRepositoryVersion(repositoryVersion2111);
    }

    // push all components to the correct version
    List<HostComponentStateEntity> hostComponentStates = hostComponentStateDAO.findByHost(hostName);
    for (HostComponentStateEntity hostComponentState : hostComponentStates) {
      hostComponentState.setVersion(repositoryVersion2111.getVersion());
      hostComponentStateDAO.merge(hostComponentState);
    }

    Map<String, String> commandParams = new HashMap<>();
    ExecutionCommand executionCommand = new ExecutionCommand();
    executionCommand.setCommandParams(commandParams);
    executionCommand.setClusterName(clusterName);

    HostRoleCommand hostRoleCommand = hostRoleCommandFactory.create(null, null, null, null);
    hostRoleCommand.setExecutionCommandWrapper(new ExecutionCommandWrapper(executionCommand));

    finalizeUpgradeAction.setExecutionCommand(executionCommand);
    finalizeUpgradeAction.setHostRoleCommand(hostRoleCommand);

    // finalize
    CommandReport report = finalizeUpgradeAction.execute(null);
    assertNotNull(report);
    assertEquals(HostRoleStatus.COMPLETED.name(), report.getStatus());

    for (HostVersionEntity hostVersion : hostVersions) {
      RepositoryVersionEntity hostRepoVersion = hostVersion.getRepositoryVersion();
      if (repositoryVersion2110.equals(hostRepoVersion)) {
        assertEquals(RepositoryVersionState.INSTALLED, hostVersion.getState());
      } else if (repositoryVersion2111.equals(hostRepoVersion)) {
        assertEquals(RepositoryVersionState.CURRENT, hostVersion.getState());
      } else {
        assertEquals(RepositoryVersionState.NOT_REQUIRED, hostVersion.getState());
      }
    }
  }

  /**
   * Tests the case where a revert happens on a patch upgrade and a new service
   * has been added which causes the repository to go OUT_OF_SYNC.
   *
   * @throws Exception
   */
  @Test
  public void testHostVersionsOutOfSyncAfterRevert() throws Exception {
    String hostName = "h1";
    Cluster cluster = createUpgradeCluster(repositoryVersion2110, hostName);
    createHostVersions(repositoryVersion2111, hostName);

    // Install ZK with some components (HBase is installed later to test the
    // logic of revert)
    Service zk = installService(cluster, "ZOOKEEPER", repositoryVersion2110);
    addServiceComponent(cluster, zk, "ZOOKEEPER_SERVER");
    addServiceComponent(cluster, zk, "ZOOKEEPER_CLIENT");
    createNewServiceComponentHost(cluster, "ZOOKEEPER", "ZOOKEEPER_SERVER", hostName);
    createNewServiceComponentHost(cluster, "ZOOKEEPER", "ZOOKEEPER_CLIENT", hostName);

    List<HostVersionEntity> hostVersions = hostVersionDAO.findAll();
    assertEquals(2, hostVersions.size());

    // repo 2110 - CURRENT
    // repo 2111 - CURRENT (PATCH)
    for (HostVersionEntity hostVersion : hostVersions) {
      hostVersion.setState(RepositoryVersionState.CURRENT);

      hostVersionDAO.merge(hostVersion);
    }

    // convert the repository into a PATCH repo (which should have ZK and HBase
    // as available services)
    repositoryVersion2111.setParent(repositoryVersion2110);
    repositoryVersion2111.setType(RepositoryType.PATCH);
    repositoryVersion2111.setVersionXml(hostName);
    repositoryVersion2111.setVersionXsd("version_definition.xsd");

    File patchVdfFile = new File("src/test/resources/hbase_version_test.xml");
    repositoryVersion2111.setVersionXml(
        IOUtils.toString(new FileInputStream(patchVdfFile), Charset.defaultCharset()));

    repositoryVersion2111 = repositoryVersionDAO.merge(repositoryVersion2111);

    // pretend like we patched
    UpgradeEntity upgrade = createUpgrade(cluster, repositoryVersion2111);
    upgrade.setOrchestration(RepositoryType.PATCH);
    upgrade.setRevertAllowed(true);
    upgrade = upgradeDAO.merge(upgrade);

    // add a service on the parent repo to cause the OUT_OF_SYNC on revert
    Service hbase = installService(cluster, "HBASE", repositoryVersion2110);
    addServiceComponent(cluster, hbase, "HBASE_MASTER");
    createNewServiceComponentHost(cluster, "HBASE", "HBASE_MASTER", hostName);

    // revert the patch
    UpgradeEntity revert = createRevert(cluster, upgrade);
    assertEquals(RepositoryType.PATCH, revert.getOrchestration());

    // push all services to the revert repo version for finalize
    Map<String, Service> services = cluster.getServices();
    assertTrue(services.size() > 0);
    for (Service service : services.values()) {
      service.setDesiredRepositoryVersion(repositoryVersion2110);
    }

    // push all components to the revert version
    List<HostComponentStateEntity> hostComponentStates = hostComponentStateDAO.findByHost(hostName);
    for (HostComponentStateEntity hostComponentState : hostComponentStates) {
      hostComponentState.setVersion(repositoryVersion2110.getVersion());
      hostComponentStateDAO.merge(hostComponentState);
    }

    Map<String, String> commandParams = new HashMap<>();
    ExecutionCommand executionCommand = new ExecutionCommand();
    executionCommand.setCommandParams(commandParams);
    executionCommand.setClusterName(clusterName);

    HostRoleCommand hostRoleCommand = hostRoleCommandFactory.create(null, null, null, null);
    hostRoleCommand.setExecutionCommandWrapper(new ExecutionCommandWrapper(executionCommand));

    finalizeUpgradeAction.setExecutionCommand(executionCommand);
    finalizeUpgradeAction.setHostRoleCommand(hostRoleCommand);

    // finalize
    CommandReport report = finalizeUpgradeAction.execute(null);
    assertNotNull(report);
    assertEquals(HostRoleStatus.COMPLETED.name(), report.getStatus());

    for (HostVersionEntity hostVersion : hostVersions) {
      RepositoryVersionEntity hostRepoVersion = hostVersion.getRepositoryVersion();
      if (repositoryVersion2110.equals(hostRepoVersion)) {
        assertEquals(RepositoryVersionState.CURRENT, hostVersion.getState());
      } else {
        assertEquals(RepositoryVersionState.OUT_OF_SYNC, hostVersion.getState());
      }
    }
  }

  private ServiceComponentHost createNewServiceComponentHost(Cluster cluster, String svc,
                                                             String svcComponent, String hostName) throws AmbariException {
    Assert.assertNotNull(cluster.getConfigGroups());
    Service s = installService(cluster, svc, sourceRepositoryVersion);
    ServiceComponent sc = addServiceComponent(cluster, s, svcComponent);

    ServiceComponentHost sch = serviceComponentHostFactory.createNew(sc, hostName);

    sc.addServiceComponentHost(sch);
    sch.setDesiredState(State.INSTALLED);
    sch.setState(State.INSTALLED);
    return sch;
  }

  private Service installService(Cluster cluster, String serviceName,
      RepositoryVersionEntity repositoryVersionEntity) throws AmbariException {
    Service service = null;

    try {
      service = cluster.getService(serviceName);
    } catch (ServiceNotFoundException e) {
      service = serviceFactory.createNew(cluster, serviceName, repositoryVersionEntity);
      cluster.addService(service);
    }

    return service;
  }

  private ServiceComponent addServiceComponent(Cluster cluster, Service service,
      String componentName) throws AmbariException {
    ServiceComponent serviceComponent = null;
    try {
      serviceComponent = service.getServiceComponent(componentName);
    } catch (ServiceComponentNotFoundException e) {
      serviceComponent = serviceComponentFactory.createNew(service, componentName);
      service.addServiceComponent(serviceComponent);
      serviceComponent.setDesiredState(State.INSTALLED);
    }

    return serviceComponent;
  }

  private void createConfigs(Cluster cluster) {
    Map<String, String> properties = new HashMap<>();
    Map<String, Map<String, String>> propertiesAttributes = new HashMap<>();
    properties.put("a", "a1");
    properties.put("b", "b1");

    configFactory.createNew(cluster, "zookeeper-env", "version-" + System.currentTimeMillis(),
        properties, propertiesAttributes);

    properties.put("zookeeper_a", "value_1");
    properties.put("zookeeper_b", "value_2");

    configFactory.createNew(cluster, "hdfs-site", "version-" + System.currentTimeMillis(),
        properties, propertiesAttributes);

    properties.put("hdfs_a", "value_3");
    properties.put("hdfs_b", "value_4");

    configFactory.createNew(cluster, "core-site", "version-" + System.currentTimeMillis(),
        properties, propertiesAttributes);

    configFactory.createNew(cluster, "foo-site", "version-" + System.currentTimeMillis(),
        properties, propertiesAttributes);
  }

  /**
   * Creates an upgrade and associates it with the cluster.
   */
  private UpgradeEntity createUpgrade(Cluster cluster, RepositoryVersionEntity repositoryVersion)
      throws Exception {

    // create some entities for the finalize action to work with for patch
    // history
    RequestEntity requestEntity = new RequestEntity();
    requestEntity.setClusterId(cluster.getClusterId());
    requestEntity.setRequestId(1L);
    requestEntity.setStartTime(System.currentTimeMillis());
    requestEntity.setCreateTime(System.currentTimeMillis());
    requestDAO.create(requestEntity);

    UpgradeEntity upgradeEntity = new UpgradeEntity();
    upgradeEntity.setId(1L);
    upgradeEntity.setClusterId(cluster.getClusterId());
    upgradeEntity.setRequestEntity(requestEntity);
    upgradeEntity.setUpgradePackage("");
    upgradeEntity.setUpgradePackStackId(new StackId((String) null));
    upgradeEntity.setRepositoryVersion(repositoryVersion);
    upgradeEntity.setUpgradeType(UpgradeType.NON_ROLLING);

    Map<String, Service> services = cluster.getServices();
    for (String serviceName : services.keySet()) {
      Service service = services.get(serviceName);
      Map<String, ServiceComponent> components = service.getServiceComponents();
      for (String componentName : components.keySet()) {
        UpgradeHistoryEntity history = new UpgradeHistoryEntity();
        history.setUpgrade(upgradeEntity);
        history.setServiceName(serviceName);
        history.setComponentName(componentName);
        history.setFromRepositoryVersion(service.getDesiredRepositoryVersion());
