/*
 * 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.state.cluster;

import static org.hamcrest.CoreMatchers.is;
import static org.junit.Assert.assertEquals;
import static org.junit.Assert.assertFalse;
import static org.junit.Assert.assertNotNull;
import static org.junit.Assert.assertThat;
import static org.junit.Assert.assertTrue;
import static org.junit.Assert.fail;
import static org.mockito.Mockito.mock;
import static org.mockito.Mockito.when;

import java.sql.SQLException;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collection;
import java.util.Collections;
import java.util.Comparator;
import java.util.ConcurrentModificationException;
import java.util.EnumSet;
import java.util.HashMap;
import java.util.HashSet;
import java.util.Iterator;
import java.util.List;
import java.util.Map;
import java.util.Set;

import jakarta.persistence.EntityManager;

import org.apache.ambari.server.AmbariException;
import org.apache.ambari.server.H2DatabaseCleaner;
import org.apache.ambari.server.agent.AgentEnv;
import org.apache.ambari.server.agent.AgentEnv.Directory;
import org.apache.ambari.server.agent.DiskInfo;
import org.apache.ambari.server.agent.HostInfo;
import org.apache.ambari.server.api.services.AmbariMetaInfo;
import org.apache.ambari.server.controller.ClusterResponse;
import org.apache.ambari.server.controller.ConfigurationResponse;
import org.apache.ambari.server.controller.ServiceConfigVersionResponse;
import org.apache.ambari.server.controller.internal.DeleteHostComponentStatusMetaData;
import org.apache.ambari.server.events.ClusterConfigChangedEvent;
import org.apache.ambari.server.events.publishers.AmbariEventPublisher;
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.ClusterDAO;
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.StackDAO;
import org.apache.ambari.server.orm.entities.ClusterConfigEntity;
import org.apache.ambari.server.orm.entities.ClusterEntity;
import org.apache.ambari.server.orm.entities.ClusterServiceEntity;
import org.apache.ambari.server.orm.entities.HostComponentStateEntity;
import org.apache.ambari.server.orm.entities.HostEntity;
import org.apache.ambari.server.orm.entities.HostStateEntity;
import org.apache.ambari.server.orm.entities.HostVersionEntity;
import org.apache.ambari.server.orm.entities.RepositoryVersionEntity;
import org.apache.ambari.server.orm.entities.ServiceDesiredStateEntity;
import org.apache.ambari.server.orm.entities.StackEntity;
import org.apache.ambari.server.state.AgentVersion;
import org.apache.ambari.server.state.Cluster;
import org.apache.ambari.server.state.Clusters;
import org.apache.ambari.server.state.ComponentInfo;
import org.apache.ambari.server.state.Config;
import org.apache.ambari.server.state.ConfigFactory;
import org.apache.ambari.server.state.ConfigHelper;
import org.apache.ambari.server.state.DesiredConfig;
import org.apache.ambari.server.state.Host;
import org.apache.ambari.server.state.HostHealthStatus;
import org.apache.ambari.server.state.HostState;
import org.apache.ambari.server.state.MaintenanceState;
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.configgroup.ConfigGroup;
import org.apache.ambari.server.state.configgroup.ConfigGroupFactory;
import org.apache.ambari.server.state.fsm.InvalidStateTransitionException;
import org.apache.ambari.server.state.host.HostHealthyHeartbeatEvent;
import org.apache.ambari.server.state.host.HostRegistrationRequestEvent;
import org.apache.ambari.server.utils.EventBusSynchronizer;
import org.apache.commons.lang.StringUtils;
import org.junit.After;
import org.junit.Before;
import org.junit.Test;

import com.google.common.collect.ImmutableList;
import com.google.common.collect.ImmutableMap;
import com.google.common.collect.Sets;
import com.google.gson.Gson;
import com.google.inject.AbstractModule;
import com.google.inject.Guice;
import com.google.inject.Injector;
import com.google.inject.persist.UnitOfWork;
import com.google.inject.util.Modules;

import junit.framework.Assert;

public class ClusterTest {

  private static final EnumSet<RepositoryVersionState> TERMINAL_VERSION_STATES =
      EnumSet.of(RepositoryVersionState.CURRENT, RepositoryVersionState.INSTALLED);

  private Clusters clusters;
  private Cluster c1;
  private Injector injector;
  private ServiceFactory serviceFactory;
  private ServiceComponentFactory serviceComponentFactory;
  private ServiceComponentHostFactory serviceComponentHostFactory;
  private AmbariMetaInfo metaInfo;
  private ConfigFactory configFactory;
  private ConfigGroupFactory configGroupFactory;
  private OrmTestHelper helper;
  private StackDAO stackDAO;
  private ClusterDAO clusterDAO;
  private HostDAO hostDAO;

  private HostVersionDAO hostVersionDAO;
  private HostComponentStateDAO hostComponentStateDAO;
  private RepositoryVersionDAO repositoryVersionDAO;
  private Gson gson;

  private static class MockModule extends AbstractModule {
    @Override
    protected void configure() {
      EventBusSynchronizer.synchronizeAmbariEventPublisher(binder());
    }
  }

  @Before
  public void setup() throws Exception {
    injector = Guice.createInjector(Modules.override(new InMemoryDefaultTestModule()).with(new MockModule()));
    injector.getInstance(GuiceJpaInitializer.class);
    EventBusSynchronizer.synchronizeAmbariEventPublisher(injector);
    clusters = injector.getInstance(Clusters.class);
    serviceFactory = injector.getInstance(ServiceFactory.class);
    configGroupFactory = injector.getInstance(ConfigGroupFactory.class);
    serviceComponentFactory = injector.getInstance(
        ServiceComponentFactory.class);
    serviceComponentHostFactory = injector.getInstance(
        ServiceComponentHostFactory.class);
    configFactory = injector.getInstance(ConfigFactory.class);
    metaInfo = injector.getInstance(AmbariMetaInfo.class);
    helper = injector.getInstance(OrmTestHelper.class);
    stackDAO = injector.getInstance(StackDAO.class);
    clusterDAO = injector.getInstance(ClusterDAO.class);
    hostDAO = injector.getInstance(HostDAO.class);
    hostVersionDAO = injector.getInstance(HostVersionDAO.class);
    hostComponentStateDAO = injector.getInstance(HostComponentStateDAO.class);
    repositoryVersionDAO = injector.getInstance(RepositoryVersionDAO.class);
    gson = injector.getInstance(Gson.class);
    injector.getInstance(UnitOfWork.class).begin();
  }

  @After
  public void teardown() throws AmbariException, SQLException {
    injector.getInstance(UnitOfWork.class).end();
    H2DatabaseCleaner.clearDatabaseAndStopPersistenceService(injector);
  }

  private void createDefaultCluster() throws Exception {
    createDefaultCluster(Sets.newHashSet("h1", "h2"));
  }

  private void createDefaultCluster(Set<String> hostNames) throws Exception {
    StackId stackId = new StackId("HDP", "0.1");
    createDefaultCluster(hostNames, stackId);
  }

  private void createDefaultCluster(Set<String> hostNames, StackId stackId) throws Exception {
    // TODO, use common function
    StackEntity stackEntity = stackDAO.find(stackId.getStackName(), stackId.getStackVersion());
    org.junit.Assert.assertNotNull(stackEntity);

    String clusterName = "c1";

    clusters.addCluster(clusterName, stackId);

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

    RepositoryVersionEntity repositoryVersion = helper.getOrCreateRepositoryVersion(stackId, stackId.getStackVersion());

    for (String hostName : hostNames) {
      clusters.addHost(hostName);

      HostEntity hostEntity = hostDAO.findByName(hostName);
      hostEntity.setIpv4("ipv4");
      hostEntity.setIpv6("ipv6");
      hostEntity.setHostAttributes(gson.toJson(hostAttributes));


      hostDAO.merge(hostEntity);
    }

    clusters.mapAndPublishHostsToCluster(hostNames, clusterName);
    c1 = clusters.getCluster(clusterName);


  }

  public ClusterEntity createDummyData() {
    StackEntity stackEntity = new StackEntity();
    stackEntity.setStackName("HDP");
    stackEntity.setStackVersion("0.1");

    ClusterEntity clusterEntity = new ClusterEntity();
    clusterEntity.setClusterId(1L);
    clusterEntity.setClusterName("test_cluster1");
    clusterEntity.setClusterInfo("test_cluster_info1");
    clusterEntity.setDesiredStack(stackEntity);

    HostEntity host1 = new HostEntity();
    HostEntity host2 = new HostEntity();
    HostEntity host3 = new HostEntity();

    host1.setHostName("test_host1");
    host2.setHostName("test_host2");
    host3.setHostName("test_host3");
    host1.setIpv4("192.168.0.1");
    host2.setIpv4("192.168.0.2");
    host3.setIpv4("192.168.0.3");

    List<HostEntity> hostEntities = new ArrayList<>();
    hostEntities.add(host1);
    hostEntities.add(host2);

    clusterEntity.setHostEntities(hostEntities);
    clusterEntity.setClusterConfigEntities(Collections.emptyList());
    //both sides of relation should be set when modifying in runtime
    host1.setClusterEntities(Arrays.asList(clusterEntity));
    host2.setClusterEntities(Arrays.asList(clusterEntity));

    HostStateEntity hostStateEntity1 = new HostStateEntity();
    hostStateEntity1.setCurrentState(HostState.HEARTBEAT_LOST);
    hostStateEntity1.setHostEntity(host1);
    HostStateEntity hostStateEntity2 = new HostStateEntity();
    hostStateEntity2.setCurrentState(HostState.HEALTHY);
    hostStateEntity2.setHostEntity(host2);
    host1.setHostStateEntity(hostStateEntity1);
    host2.setHostStateEntity(hostStateEntity2);

    ClusterServiceEntity clusterServiceEntity = new ClusterServiceEntity();
    clusterServiceEntity.setServiceName("HDFS");
    clusterServiceEntity.setClusterEntity(clusterEntity);
    clusterServiceEntity.setServiceComponentDesiredStateEntities(
        Collections.emptyList());

    ServiceDesiredStateEntity stateEntity = mock(ServiceDesiredStateEntity.class);

    when(stateEntity.getDesiredStack()).thenReturn(stackEntity);

    clusterServiceEntity.setServiceDesiredStateEntity(stateEntity);
    List<ClusterServiceEntity> clusterServiceEntities = new ArrayList<>();
    clusterServiceEntities.add(clusterServiceEntity);
    clusterEntity.setClusterServiceEntities(clusterServiceEntities);
    return clusterEntity;
  }

  /**
   * For Rolling Upgrades, create a cluster with the following components HDFS:
   * NameNode, DataNode, HDFS Client ZK: Zookeeper Server, Zookeeper Monitor
   * Ganglia: Ganglia Server, Ganglia Monitor
   *
   * Further, 3 hosts will be added. Finally, verify that only the Ganglia
   * components do not need to advertise a version.
   *
   * @param clusterName
   *          Cluster Name
   * @param repositoryVersion
   *          the repository to use for new services being installed in the
   *          cluster
   * @param hostAttributes
   *          Host attributes to use for 3 hosts (h-1, h-2, h-3)
   * @throws Exception
   * @return Cluster that was created
   */
  private Cluster createClusterForRU(String clusterName, RepositoryVersionEntity repositoryVersion,
      Map<String, String> hostAttributes) throws Exception {
    StackId stackId = repositoryVersion.getStackId();
    clusters.addCluster(clusterName, stackId);
    Cluster cluster = clusters.getCluster(clusterName);
    Assert.assertEquals(clusterName, cluster.getClusterName());
    //Assert.assertEquals(1, cluster.getClusterId());

    // Add Hosts
    List<String> hostNames = new ArrayList<String>() {{ add("h-1"); add("h-2"); add("h-3"); }};
    for(String hostName : hostNames) {
      addHost(hostName, hostAttributes);
    }

    // Add stack and map Hosts to cluster
    cluster.setDesiredStackVersion(stackId);
    cluster.setCurrentStackVersion(stackId);
    for(String hostName : hostNames) {
      clusters.mapHostToCluster(hostName, clusterName);
    }

    // Transition all hosts to HEALTHY state
    for (Host host : cluster.getHosts()) {
      host.setState(HostState.HEALTHY);
    }

    // Add Services
    Service s1 = serviceFactory.createNew(cluster, "HDFS", repositoryVersion);
    Service s2 = serviceFactory.createNew(cluster, "ZOOKEEPER", repositoryVersion);
    Service s3 = serviceFactory.createNew(cluster, "GANGLIA", repositoryVersion);
    cluster.addService(s1);
    cluster.addService(s2);
    cluster.addService(s3);

    // Add HDFS components
    ServiceComponent sc1CompA = serviceComponentFactory.createNew(s1, "NAMENODE");
    ServiceComponent sc1CompB = serviceComponentFactory.createNew(s1, "DATANODE");
    ServiceComponent sc1CompC = serviceComponentFactory.createNew(s1, "HDFS_CLIENT");
    s1.addServiceComponent(sc1CompA);
    s1.addServiceComponent(sc1CompB);
    s1.addServiceComponent(sc1CompC);

    // Add ZK
    ServiceComponent sc2CompA = serviceComponentFactory.createNew(s2, "ZOOKEEPER_SERVER");
    ServiceComponent sc2CompB = serviceComponentFactory.createNew(s2, "ZOOKEEPER_CLIENT");
    s2.addServiceComponent(sc2CompA);
    s2.addServiceComponent(sc2CompB);

    // Add Ganglia
    ServiceComponent sc3CompA = serviceComponentFactory.createNew(s3, "GANGLIA_SERVER");
    ServiceComponent sc3CompB = serviceComponentFactory.createNew(s3, "GANGLIA_MONITOR");
    s3.addServiceComponent(sc3CompA);
    s3.addServiceComponent(sc3CompB);

    // Host 1 will have all components
    ServiceComponentHost schHost1Serv1CompA = serviceComponentHostFactory.createNew(sc1CompA, "h-1");
    ServiceComponentHost schHost1Serv1CompB = serviceComponentHostFactory.createNew(sc1CompB, "h-1");
    ServiceComponentHost schHost1Serv1CompC = serviceComponentHostFactory.createNew(sc1CompC, "h-1");
    ServiceComponentHost schHost1Serv2CompA = serviceComponentHostFactory.createNew(sc2CompA, "h-1");
    ServiceComponentHost schHost1Serv2CompB = serviceComponentHostFactory.createNew(sc2CompB, "h-1");
    ServiceComponentHost schHost1Serv3CompA = serviceComponentHostFactory.createNew(sc3CompA, "h-1");
    ServiceComponentHost schHost1Serv3CompB = serviceComponentHostFactory.createNew(sc3CompB, "h-1");
    sc1CompA.addServiceComponentHost(schHost1Serv1CompA);
    sc1CompB.addServiceComponentHost(schHost1Serv1CompB);
    sc1CompC.addServiceComponentHost(schHost1Serv1CompC);
    sc2CompA.addServiceComponentHost(schHost1Serv2CompA);
    sc2CompB.addServiceComponentHost(schHost1Serv2CompB);
    sc3CompA.addServiceComponentHost(schHost1Serv3CompA);
    sc3CompB.addServiceComponentHost(schHost1Serv3CompB);

    // Host 2 will have ZK_CLIENT and GANGLIA_MONITOR
    ServiceComponentHost schHost2Serv2CompB = serviceComponentHostFactory.createNew(sc2CompB, "h-2");
    ServiceComponentHost schHost2Serv3CompB = serviceComponentHostFactory.createNew(sc3CompB, "h-2");
    sc2CompB.addServiceComponentHost(schHost2Serv2CompB);
    sc3CompB.addServiceComponentHost(schHost2Serv3CompB);

    // Host 3 will have GANGLIA_MONITOR
    ServiceComponentHost schHost3Serv3CompB = serviceComponentHostFactory.createNew(sc3CompB, "h-3");
    sc3CompB.addServiceComponentHost(schHost3Serv3CompB);

    // Verify count of components
    List<ServiceComponentHost> scHost1 = cluster.getServiceComponentHosts("h-1");
    Assert.assertEquals(7, scHost1.size());

    List<ServiceComponentHost> scHost2 = cluster.getServiceComponentHosts("h-2");
    Assert.assertEquals(2, scHost2.size());

    List<ServiceComponentHost> scHost3 = cluster.getServiceComponentHosts("h-3");
    Assert.assertEquals(1, scHost3.size());

    //<editor-fold desc="Validate Version Advertised">
    /*
    For some reason this still uses the metainfo.xml files for these services
    from HDP-2.0.5 stack instead of the provided Stack Id
    */
    HashMap<String, Set<String>> componentsThatAdvertiseVersion = new HashMap<>();
    HashMap<String, Set<String>> componentsThatDontAdvertiseVersion = new HashMap<>();

    Set<String> hdfsComponents = new HashSet<String>() {{ add("NAMENODE"); add("DATANODE"); add("HDFS_CLIENT"); }};
    Set<String> zkComponents = new HashSet<String>() {{ add("ZOOKEEPER_SERVER"); add("ZOOKEEPER_CLIENT"); }};
    Set<String> gangliaComponents = new HashSet<String>() {{ add("GANGLIA_SERVER"); add("GANGLIA_MONITOR"); }};

    componentsThatAdvertiseVersion.put("HDFS", hdfsComponents);
    componentsThatAdvertiseVersion.put("ZOOKEEPER", zkComponents);
    componentsThatDontAdvertiseVersion.put("GANGLIA", gangliaComponents);

    for(String service : componentsThatAdvertiseVersion.keySet())  {
      Set<String> components = componentsThatAdvertiseVersion.get(service);
      for(String componentName : components) {
        ComponentInfo component = metaInfo.getComponent(stackId.getStackName(), stackId.getStackVersion(), service, componentName);
        Assert.assertTrue(component.isVersionAdvertised());
      }
    }

    for(String service : componentsThatDontAdvertiseVersion.keySet())  {
      Set<String> components = componentsThatDontAdvertiseVersion.get(service);
      for(String componentName : components) {
        ComponentInfo component = metaInfo.getComponent(stackId.getStackName(), stackId.getStackVersion(), service, componentName);
        Assert.assertFalse(component.isVersionAdvertised());
      }
    }
    //</editor-fold>

    return cluster;
  }

  /**
   * Add a host to the system with the provided attributes.
   * @param hostName Host Name
   * @param hostAttributes Host Attributes
   * @throws Exception
   */
  private Host addHost(String hostName, Map<String, String> hostAttributes) throws Exception {
    clusters.addHost(hostName);
    Host host = clusters.getHost(hostName);
    host.setIPv4("ipv4");
    host.setIPv6("ipv6");
    host.setHostAttributes(hostAttributes);

    return host;
  }

  /**
   * For the provided collection of HostComponentStates, set the version to {@paramref version} if the Component
   * can advertise a version. Then, simulate the {@link org.apache.ambari.server.events.listeners.upgrade.StackVersionListener}
   * by calling methods to transition the HostVersion, and recalculate the ClusterVersion.
   * @param stackId Stack ID to retrieve the ComponentInfo
   * @param version Version to set
   * @param cluster Cluster to retrieve services from
   * @param hostComponentStates Collection to set the version for
   * @throws Exception
   */
  private void simulateStackVersionListener(StackId stackId, String version, Cluster cluster, List<HostComponentStateEntity> hostComponentStates) throws Exception {
    for(int i = 0; i < hostComponentStates.size(); i++) {
      HostComponentStateEntity hce = hostComponentStates.get(i);
      ComponentInfo compInfo = metaInfo.getComponent(
          stackId.getStackName(), stackId.getStackVersion(),
          hce.getServiceName(),
          hce.getComponentName());

      if (compInfo.isVersionAdvertised()) {
        hce.setVersion(version);
        hostComponentStateDAO.merge(hce);
      }

      helper.getOrCreateRepositoryVersion(stackId, version);

      // Simulate the StackVersionListener during the installation
      Service svc = cluster.getService(hce.getServiceName());
      ServiceComponent svcComp = svc.getServiceComponent(hce.getComponentName());
      ServiceComponentHost scHost = svcComp.getServiceComponentHost(hce.getHostName());

      scHost.recalculateHostVersionState();
    }
  }

  @Test
  public void testAddHost() throws Exception {
    createDefaultCluster();
    clusters.addHost("h3");

    try {
      clusters.addHost("h3");
      fail("Duplicate add should fail");
    }
    catch (AmbariException e) {
      // Expected
    }
  }

  @Test
  public void testGetHostState() throws Exception {
    createDefaultCluster();

    Assert.assertEquals(HostState.INIT, clusters.getHost("h1").getState());
  }

  @Test
  public void testSetHostState() throws Exception {
    createDefaultCluster();

    clusters.getHost("h1").setState(HostState.HEARTBEAT_LOST);
    Assert.assertEquals(HostState.HEARTBEAT_LOST,
        clusters.getHost("h1").getState());
  }

  @Test
  public void testHostEvent() throws Exception,
      InvalidStateTransitionException {
    createDefaultCluster();

    HostInfo hostInfo = new HostInfo();
    hostInfo.setHostName("h1");
    hostInfo.setInterfaces("fip_4");
    hostInfo.setArchitecture("os_arch");
    hostInfo.setOS("os_type");
    hostInfo.setMemoryTotal(10);
    hostInfo.setMemorySize(100);
    hostInfo.setProcessorCount(10);
    List<DiskInfo> mounts = new ArrayList<>();
    mounts.add(new DiskInfo("/dev/sda", "/mnt/disk1",
        "5000000", "4000000", "10%", "size", "fstype"));
    hostInfo.setMounts(mounts);

    AgentEnv agentEnv = new AgentEnv();

    Directory dir1 = new Directory();
    dir1.setName("/etc/hadoop");
    dir1.setType("not_exist");
    Directory dir2 = new Directory();
    dir2.setName("/var/log/hadoop");
    dir2.setType("not_exist");
    agentEnv.setStackFoldersAndFiles(new Directory[] { dir1, dir2 });

    AgentVersion agentVersion = new AgentVersion("0.0.x");
    long currentTime = 1001;

    clusters.getHost("h1").handleEvent(new HostRegistrationRequestEvent(
        "h1", agentVersion, currentTime, hostInfo, agentEnv, currentTime));

    Assert.assertEquals(HostState.WAITING_FOR_HOST_STATUS_UPDATES,
        clusters.getHost("h1").getState());

    clusters.getHost("h1").setState(HostState.HEARTBEAT_LOST);

    try {
      clusters.getHost("h1").handleEvent(
          new HostHealthyHeartbeatEvent("h1", currentTime, null, null));
      fail("Exception should be thrown on invalid event");
    }
    catch (InvalidStateTransitionException e) {
      // Expected
    }

  }

  @Test
  public void testBasicClusterSetup() throws Exception {
    StackId stackVersion = new StackId("HDP-1.2.0");

    createDefaultCluster();

    String clusterName = "c2";

    try {
      clusters.getCluster(clusterName);
      fail("Exception expected for invalid cluster");
    } catch (Exception e) {
      // Expected
    }

    clusters.addCluster(clusterName, stackVersion);
    Cluster c2 = clusters.getCluster(clusterName);

    Assert.assertNotNull(c2);

    Assert.assertEquals(clusterName, c2.getClusterName());

    c2.setClusterName("foo2");
    Assert.assertEquals("foo2", c2.getClusterName());

    Assert.assertNotNull(c2.getDesiredStackVersion());
    Assert.assertEquals("HDP-1.2.0", c2.getDesiredStackVersion().getStackId());
  }

  @Test
  public void testAddAndGetServices() throws Exception {
    createDefaultCluster();

    // TODO write unit tests for
    // public void addService(Service service) throws AmbariException;
    // public Service getService(String serviceName) throws AmbariException;
    // public Map<String, Service> getServices();

    RepositoryVersionEntity repositoryVersion = helper.getOrCreateRepositoryVersion(c1);

    serviceFactory.createNew(c1, "HDFS", repositoryVersion);
    serviceFactory.createNew(c1, "MAPREDUCE", repositoryVersion);

    Service s = c1.getService("HDFS");
    Assert.assertNotNull(s);
    Assert.assertEquals("HDFS", s.getName());
    Assert.assertEquals(c1.getClusterId(), s.getClusterId());

    try {
      c1.getService("HBASE");
      fail("Expected error for unknown service");
    } catch (Exception e) {
      // Expected
    }

    Map<String, Service> services = c1.getServices();
    Assert.assertEquals(2, services.size());
    Assert.assertTrue(services.containsKey("HDFS"));
    Assert.assertTrue(services.containsKey("MAPREDUCE"));
  }

  @Test
  public void testGetServiceComponentHosts() throws Exception {
    createDefaultCluster();

    // TODO write unit tests
    // public List<ServiceComponentHost> getServiceComponentHosts(String hostname);

    RepositoryVersionEntity repositoryVersion = helper.getOrCreateRepositoryVersion(c1);

    Service s = serviceFactory.createNew(c1, "HDFS", repositoryVersion);
    c1.addService(s);
    ServiceComponent sc = serviceComponentFactory.createNew(s, "NAMENODE");
    s.addServiceComponent(sc);

    ServiceComponentHost sch =
        serviceComponentHostFactory.createNew(sc, "h1");
    sc.addServiceComponentHost(sch);

    List<ServiceComponentHost> scHosts = c1.getServiceComponentHosts("h1");
    Assert.assertEquals(1, scHosts.size());

    Iterator<ServiceComponentHost> iterator = scHosts.iterator();

    //Try to iterate on sch and modify it in loop
    try {
      while (iterator.hasNext()) {
        iterator.next();
        Service s1 = serviceFactory.createNew(c1, "PIG", repositoryVersion);
        c1.addService(s1);
        ServiceComponent sc1 = serviceComponentFactory.createNew(s1, "PIG");
        s1.addServiceComponent(sc1);
        ServiceComponentHost sch1 = serviceComponentHostFactory.createNew(sc1, "h1");
        sc1.addServiceComponentHost(sch1);
      }
    } catch (ConcurrentModificationException e ) {
      Assert.assertTrue("Failed to work concurrently with sch", false);
    }

    scHosts = c1.getServiceComponentHosts("h1");
    Assert.assertEquals(2, scHosts.size());
  }

  @Test
  public void testGetServiceComponentHosts_ForService() throws Exception {
    createDefaultCluster();

    RepositoryVersionEntity repositoryVersion = helper.getOrCreateRepositoryVersion(c1);

    Service s = serviceFactory.createNew(c1, "HDFS", repositoryVersion);
    c1.addService(s);

    ServiceComponent scNN = serviceComponentFactory.createNew(s, "NAMENODE");
    s.addServiceComponent(scNN);
    ServiceComponentHost schNNH1 = serviceComponentHostFactory.createNew(scNN, "h1");
    scNN.addServiceComponentHost(schNNH1);

    ServiceComponent scDN = serviceComponentFactory.createNew(s, "DATANODE");
    s.addServiceComponent(scDN);
    ServiceComponentHost scDNH1 = serviceComponentHostFactory.createNew(scDN, "h1");
    scDN.addServiceComponentHost(scDNH1);
    ServiceComponentHost scDNH2 = serviceComponentHostFactory.createNew(scDN, "h2");
    scDN.addServiceComponentHost(scDNH2);

    List<ServiceComponentHost> scHosts;

    scHosts = c1.getServiceComponentHosts("HDFS", null);
    Assert.assertEquals(3, scHosts.size());

    scHosts = c1.getServiceComponentHosts("UNKNOWN SERVICE", null);
    Assert.assertEquals(0, scHosts.size());
  }

  @Test
  public void testGetServiceComponentHosts_ForServiceComponent() throws Exception {
    createDefaultCluster();

    RepositoryVersionEntity repositoryVersion = helper.getOrCreateRepositoryVersion(c1);

    Service s = serviceFactory.createNew(c1, "HDFS", repositoryVersion);
    c1.addService(s);

    ServiceComponent scNN = serviceComponentFactory.createNew(s, "NAMENODE");
    s.addServiceComponent(scNN);
    ServiceComponentHost schNNH1 = serviceComponentHostFactory.createNew(scNN, "h1");
    scNN.addServiceComponentHost(schNNH1);

    ServiceComponent scDN = serviceComponentFactory.createNew(s, "DATANODE");
    s.addServiceComponent(scDN);
    ServiceComponentHost scDNH1 = serviceComponentHostFactory.createNew(scDN, "h1");
    scDN.addServiceComponentHost(scDNH1);
    ServiceComponentHost scDNH2 = serviceComponentHostFactory.createNew(scDN, "h2");
    scDN.addServiceComponentHost(scDNH2);

    List<ServiceComponentHost> scHosts;

    scHosts = c1.getServiceComponentHosts("HDFS", "DATANODE");
    Assert.assertEquals(2, scHosts.size());

    scHosts = c1.getServiceComponentHosts("HDFS", "UNKNOWN COMPONENT");
    Assert.assertEquals(0, scHosts.size());

    scHosts = c1.getServiceComponentHosts("UNKNOWN SERVICE", "DATANODE");
    Assert.assertEquals(0, scHosts.size());

    scHosts = c1.getServiceComponentHosts("UNKNOWN SERVICE", "UNKNOWN COMPONENT");
    Assert.assertEquals(0, scHosts.size());
  }

  @Test
  public void testGetServiceComponentHostMap() throws Exception {
    createDefaultCluster();

    RepositoryVersionEntity repositoryVersion = helper.getOrCreateRepositoryVersion(c1);

    Service s = serviceFactory.createNew(c1, "HDFS", repositoryVersion);
    c1.addService(s);

    ServiceComponent scNN = serviceComponentFactory.createNew(s, "NAMENODE");
    s.addServiceComponent(scNN);
    ServiceComponentHost schNNH1 = serviceComponentHostFactory.createNew(scNN, "h1");
    scNN.addServiceComponentHost(schNNH1);

    ServiceComponent scDN = serviceComponentFactory.createNew(s, "DATANODE");
    s.addServiceComponent(scDN);
    ServiceComponentHost scDNH1 = serviceComponentHostFactory.createNew(scDN, "h1");
    scDN.addServiceComponentHost(scDNH1);
    ServiceComponentHost scDNH2 = serviceComponentHostFactory.createNew(scDN, "h2");
    scDN.addServiceComponentHost(scDNH2);

    Map<String, Set<String>> componentHostMap;

    componentHostMap = c1.getServiceComponentHostMap(null, null);
    Assert.assertEquals(2, componentHostMap.size());

    Assert.assertEquals(1, componentHostMap.get("NAMENODE").size());
    Assert.assertTrue(componentHostMap.get("NAMENODE").contains("h1"));

    Assert.assertEquals(2, componentHostMap.get("DATANODE").size());
    Assert.assertTrue(componentHostMap.get("DATANODE").contains("h1"));
    Assert.assertTrue(componentHostMap.get("DATANODE").contains("h2"));
  }

  @Test
  public void testGetServiceComponentHostMap_ForService() throws Exception {
    createDefaultCluster();

    RepositoryVersionEntity repositoryVersion = helper.getOrCreateRepositoryVersion(c1);

    Service sfHDFS = serviceFactory.createNew(c1, "HDFS", repositoryVersion);
    c1.addService(sfHDFS);

    Service sfMR = serviceFactory.createNew(c1, "MAPREDUCE", repositoryVersion);
    c1.addService(sfMR);

    ServiceComponent scNN = serviceComponentFactory.createNew(sfHDFS, "NAMENODE");
    sfHDFS.addServiceComponent(scNN);
    ServiceComponentHost schNNH1 = serviceComponentHostFactory.createNew(scNN, "h1");
    scNN.addServiceComponentHost(schNNH1);

    ServiceComponent scDN = serviceComponentFactory.createNew(sfHDFS, "DATANODE");
    sfHDFS.addServiceComponent(scDN);
    ServiceComponentHost scDNH1 = serviceComponentHostFactory.createNew(scDN, "h1");
    scDN.addServiceComponentHost(scDNH1);
    ServiceComponentHost scDNH2 = serviceComponentHostFactory.createNew(scDN, "h2");
    scDN.addServiceComponentHost(scDNH2);

    ServiceComponent scJT = serviceComponentFactory.createNew(sfMR, "JOBTRACKER");
    sfMR.addServiceComponent(scJT);
    ServiceComponentHost schJTH1 = serviceComponentHostFactory.createNew(scJT, "h1");
    scJT.addServiceComponentHost(schJTH1);

    Map<String, Set<String>> componentHostMap;

    componentHostMap = c1.getServiceComponentHostMap(null, Collections.singleton("HDFS"));
    Assert.assertEquals(2, componentHostMap.size());
    Assert.assertEquals(1, componentHostMap.get("NAMENODE").size());
    Assert.assertTrue(componentHostMap.get("NAMENODE").contains("h1"));
    Assert.assertEquals(2, componentHostMap.get("DATANODE").size());
    Assert.assertTrue(componentHostMap.get("DATANODE").contains("h1"));
    Assert.assertTrue(componentHostMap.get("DATANODE").contains("h2"));

    componentHostMap = c1.getServiceComponentHostMap(null, Collections.singleton("MAPREDUCE"));
    Assert.assertEquals(1, componentHostMap.size());
    Assert.assertEquals(1, componentHostMap.get("JOBTRACKER").size());
    Assert.assertTrue(componentHostMap.get("JOBTRACKER").contains("h1"));

