content: /*
   * 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.host;
  
  import static org.junit.Assert.fail;
  import static org.mockito.Matchers.any;
  import static org.mockito.Mockito.doNothing;
  import static org.mockito.Mockito.mock;
  import static org.mockito.Mockito.when;
  
  import java.sql.SQLException;
  import java.util.ArrayList;
  import java.util.HashMap;
  import java.util.List;
  import java.util.Map;
  
  import org.apache.ambari.server.AmbariException;
  import org.apache.ambari.server.H2DatabaseCleaner;
  import org.apache.ambari.server.actionmanager.ActionManager;
  import org.apache.ambari.server.agent.AgentEnv;
  import org.apache.ambari.server.agent.DiskInfo;
  import org.apache.ambari.server.agent.HeartBeatHandler;
  import org.apache.ambari.server.agent.HostInfo;
  import org.apache.ambari.server.api.services.AmbariMetaInfo;
  import org.apache.ambari.server.configuration.Configuration;
  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.HostDAO;
  import org.apache.ambari.server.orm.entities.HostEntity;
  import org.apache.ambari.server.orm.entities.HostStateEntity;
  import org.apache.ambari.server.security.encryption.Encryptor;
  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.Config;
  import org.apache.ambari.server.state.ConfigFactory;
  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.HostHealthStatus.HealthStatus;
  import org.apache.ambari.server.state.HostState;
  import org.apache.ambari.server.state.MaintenanceState;
  import org.apache.ambari.server.state.StackId;
  import org.junit.After;
  import org.junit.Assert;
  import org.junit.Before;
  import org.junit.Test;
  import org.slf4j.Logger;
  import org.slf4j.LoggerFactory;
  
  import com.google.inject.Guice;
  import com.google.inject.Injector;
  
  public class HostTest {
  
    private Injector injector;
    private Clusters clusters;
    private HostDAO hostDAO;
    private OrmTestHelper helper;
    private static final Logger LOG = LoggerFactory.getLogger(HostTest.class);
  
    @Before
     public void setup() throws AmbariException{
      injector = Guice.createInjector(new InMemoryDefaultTestModule());
      injector.getInstance(GuiceJpaInitializer.class);
      clusters = injector.getInstance(Clusters.class);
      hostDAO = injector.getInstance(HostDAO.class);
      helper = injector.getInstance(OrmTestHelper.class);
    }
  
    @After
    public void teardown() throws AmbariException, SQLException {
      H2DatabaseCleaner.clearDatabaseAndStopPersistenceService(injector);
    }
  
    @Test
    public void testHostInfoImport() throws AmbariException{
      HostInfo info = new HostInfo();
      info.setMemorySize(100);
      info.setProcessorCount(10);
      info.setPhysicalProcessorCount(2);
      List<DiskInfo> mounts = new ArrayList<>();
      mounts.add(new DiskInfo("/dev/sda", "/mnt/disk1",
          "5000000", "4000000", "10%", "size", "fstype"));
      info.setMounts(mounts);
  
      info.setHostName("foo");
      info.setInterfaces("fip_4");
      info.setArchitecture("os_arch");
      info.setOS("os_type");
      info.setMemoryTotal(10);
  
      clusters.addHost("foo");
      Host host = clusters.getHost("foo");
  
      host.importHostInfo(info);
  
      Assert.assertEquals(info.getHostName(), host.getHostName());
      Assert.assertEquals(info.getFreeMemory(), host.getAvailableMemBytes());
      Assert.assertEquals(info.getMemoryTotal(), host.getTotalMemBytes());
      Assert.assertEquals(info.getProcessorCount(), host.getCpuCount());
      Assert.assertEquals(info.getPhysicalProcessorCount(), host.getPhCpuCount());
      Assert.assertEquals(info.getMounts().size(), host.getDisksInfo().size());
      Assert.assertEquals(info.getArchitecture(), host.getOsArch());
      Assert.assertEquals(info.getOS(), host.getOsType());
    }
  
    private void registerHost(Host host) throws Exception {
      registerHost(host, true);
    }
  
    @Test
    public void testHostOs() throws Exception {
      Clusters clusters = mock(Clusters.class);
      ActionManager manager = mock(ActionManager.class);
      Injector injector = mock(Injector.class);
      Configuration configuration = mock(Configuration.class);
      doNothing().when(injector).injectMembers(any());
      when(injector.getInstance(AmbariEventPublisher.class)).thenReturn(mock(AmbariEventPublisher.class));
      HeartBeatHandler handler = new HeartBeatHandler(configuration, clusters, manager, Encryptor.NONE, injector);
      String os = handler.getOsType("RedHat", "6.1");
      Assert.assertEquals("redhat6", os);
      os = handler.getOsType("RedHat", "6");
      Assert.assertEquals("redhat6", os);
      os = handler.getOsType("RedHat6","");
      Assert.assertEquals("redhat6", os);
  
    }
  
    private void registerHost(Host host, boolean firstReg) throws Exception {
      HostInfo info = new HostInfo();
      info.setMemorySize(100);
      info.setProcessorCount(10);
      List<DiskInfo> mounts = new ArrayList<>();
      mounts.add(new DiskInfo("/dev/sda", "/mnt/disk1",
          "5000000", "4000000", "10%", "size", "fstype"));
      info.setMounts(mounts);
  
      info.setHostName("foo");
      info.setInterfaces("fip_4");
      info.setArchitecture("os_arch");
      info.setOS("os_type");
      info.setMemoryTotal(10);
  
      AgentVersion agentVersion = null;
      long currentTime = System.currentTimeMillis();
  
      AgentEnv agentEnv = new AgentEnv();
  
      HostRegistrationRequestEvent e =
          new HostRegistrationRequestEvent("foo", agentVersion, currentTime,
              info, agentEnv, currentTime);
  
      if (!firstReg) {
        Assert.assertNotNull(host.getHostId());
      }
  
      host.handleEvent(e);
      Assert.assertEquals(currentTime, host.getLastRegistrationTime());
  
      Assert.assertNotNull(host.getLastAgentEnv());
  
      HostEntity entity = hostDAO.findByName(host.getHostName());
      Assert.assertEquals(currentTime,
          entity.getLastRegistrationTime().longValue());
      Assert.assertEquals("os_arch", entity.getOsArch());
      Assert.assertEquals("os_type", entity.getOsType());
      Assert.assertEquals(10, entity.getTotalMem().longValue());
    }
  
    private void ensureHostUpdatesReceived(Host host) throws Exception {
      HostStatusUpdatesReceivedEvent e =
          new HostStatusUpdatesReceivedEvent(host.getHostName(), 1);
      host.handleEvent(e);
    }
  
    private void verifyHostState(Host host, HostState state) {
      Assert.assertEquals(state, host.getState());
    }
  
    private void sendHealthyHeartbeat(Host host, long counter)
        throws Exception {
      HostHealthyHeartbeatEvent e = new HostHealthyHeartbeatEvent(
          host.getHostName(), counter, null, null);
      host.handleEvent(e);
    }
  
    private void sendUnhealthyHeartbeat(Host host, long counter)
        throws Exception {
      HostHealthStatus healthStatus = new HostHealthStatus(HealthStatus.UNHEALTHY,
          "Unhealthy server");
      HostUnhealthyHeartbeatEvent e = new HostUnhealthyHeartbeatEvent(
          host.getHostName(), counter, healthStatus);
      host.handleEvent(e);
    }
  
    private void timeoutHost(Host host) throws Exception {
      HostHeartbeatLostEvent e = new HostHeartbeatLostEvent(
          host.getHostName());
      host.handleEvent(e);
    }
  
    @Test
    public void testHostFSMInit() throws AmbariException{
      clusters.addHost("foo");
      Host host = clusters.getHost("foo");
      verifyHostState(host, HostState.INIT);
    }
  
    @Test
    public void testHostRegistrationFlow() throws Exception {
      clusters.addHost("foo");
      Host host = clusters.getHost("foo");
      registerHost(host);
      verifyHostState(host, HostState.WAITING_FOR_HOST_STATUS_UPDATES);
  
      boolean exceptionThrown = false;
      try {
        registerHost(host);
      } catch (Exception e) {
        // Expected
        exceptionThrown = true;
      }
      if (!exceptionThrown) {
        fail("Expected invalid transition exception to be thrown");
      }
  
      ensureHostUpdatesReceived(host);
      verifyHostState(host, HostState.HEALTHY);
  
      exceptionThrown = false;
      try {
        ensureHostUpdatesReceived(host);
      } catch (Exception e) {
        // Expected
        exceptionThrown = true;
      }
      if (!exceptionThrown) {
        fail("Expected invalid transition exception to be thrown");
      }
    }
  
    @Test
    public void testHostHeartbeatFlow() throws Exception {
      clusters.addHost("foo");
      Host host = clusters.getHost("foo");
      registerHost(host);
      ensureHostUpdatesReceived(host);
  
      // TODO need to verify audit logs generated
      // TODO need to verify health status updated properly
  
      long counter = 0;
      sendHealthyHeartbeat(host, ++counter);
      verifyHostState(host, HostState.HEALTHY);
      Assert.assertEquals(counter, host.getLastHeartbeatTime());
  
      sendHealthyHeartbeat(host, ++counter);
      verifyHostState(host, HostState.HEALTHY);
      Assert.assertEquals(counter, host.getLastHeartbeatTime());
      Assert.assertEquals(HealthStatus.HEALTHY,
          host.getHealthStatus().getHealthStatus());
  
      sendUnhealthyHeartbeat(host, ++counter);
      verifyHostState(host, HostState.UNHEALTHY);
      Assert.assertEquals(counter, host.getLastHeartbeatTime());
      Assert.assertEquals(HealthStatus.UNHEALTHY,
          host.getHealthStatus().getHealthStatus());
  
      sendUnhealthyHeartbeat(host, ++counter);
      verifyHostState(host, HostState.UNHEALTHY);
      Assert.assertEquals(counter, host.getLastHeartbeatTime());
      Assert.assertEquals(HealthStatus.UNHEALTHY,
          host.getHealthStatus().getHealthStatus());
  
      sendHealthyHeartbeat(host, ++counter);
      verifyHostState(host, HostState.HEALTHY);
      Assert.assertEquals(counter, host.getLastHeartbeatTime());
      Assert.assertEquals(HealthStatus.HEALTHY,
          host.getHealthStatus().getHealthStatus());
  
      timeoutHost(host);
      verifyHostState(host, HostState.HEARTBEAT_LOST);
      Assert.assertEquals(counter, host.getLastHeartbeatTime());
      Assert.assertEquals(HealthStatus.UNKNOWN,
          host.getHealthStatus().getHealthStatus());
  
      timeoutHost(host);
      verifyHostState(host, HostState.HEARTBEAT_LOST);
      Assert.assertEquals(counter, host.getLastHeartbeatTime());
      Assert.assertEquals(HealthStatus.UNKNOWN,
          host.getHealthStatus().getHealthStatus());
  
      try {
        sendUnhealthyHeartbeat(host, ++counter);
        fail("Invalid event should have triggered an exception");
      } catch (Exception e) {
        // Expected
      }
      verifyHostState(host, HostState.HEARTBEAT_LOST);
  
      try {
        sendHealthyHeartbeat(host, ++counter);
        fail("Invalid event should have triggered an exception");
      } catch (Exception e) {
        // Expected
      }
      verifyHostState(host, HostState.HEARTBEAT_LOST);
    }
  
    @Test
    public void testHostRegistrationsInAnyState() throws Exception {
      clusters.addHost("foo");
      Host host = clusters.getHost("foo");
      host.setIPv4("ipv4");
      host.setIPv6("ipv6");
  
      long counter = 0;
  
      registerHost(host);
  
      ensureHostUpdatesReceived(host);
      registerHost(host, false);
  
      ensureHostUpdatesReceived(host);
      sendHealthyHeartbeat(host, ++counter);
      verifyHostState(host, HostState.HEALTHY);
      registerHost(host, false);
      ensureHostUpdatesReceived(host);
  
      sendUnhealthyHeartbeat(host, ++counter);
      verifyHostState(host, HostState.UNHEALTHY);
      registerHost(host, false);
      ensureHostUpdatesReceived(host);
  
      timeoutHost(host);
      verifyHostState(host, HostState.HEARTBEAT_LOST);
      registerHost(host, false);
      ensureHostUpdatesReceived(host);
  
      host.setState(HostState.INIT);
      registerHost(host, false);
  
    }
  
    @Test
    public void testHostDesiredConfig() throws Exception {
      AmbariMetaInfo metaInfo = injector.getInstance(AmbariMetaInfo.class);
  
      StackId stackId = new StackId("HDP-0.1");
      clusters.addCluster("c1", stackId);
      Cluster c1 = clusters.getCluster("c1");
  
      helper.getOrCreateRepositoryVersion(stackId, stackId.getStackVersion());
  
      Assert.assertEquals("c1", c1.getClusterName());
      clusters.addHost("h1");
      Host host = clusters.getHost("h1");
      host.setIPv4("ipv4");
      host.setIPv6("ipv6");
  
      Map<String, String> hostAttributes = new HashMap<>();
      hostAttributes.put("os_family", "redhat");
      hostAttributes.put("os_release_version", "6.3");
      host.setHostAttributes(hostAttributes);
  
      c1.setDesiredStackVersion(stackId);
      clusters.mapHostToCluster("h1", "c1");
  
      ConfigFactory configFactory = injector.getInstance(ConfigFactory.class);
      Config config = configFactory.createNew(c1, "global", "v1",
          new HashMap<String,String>() {{ put("a", "b"); put("x", "y"); }}, new HashMap<>());
  
      try {
        host.addDesiredConfig(c1.getClusterId(), true, null, config);
        Assert.fail("Expect failure when user is not specified.");
      }
      catch (Exception e) {
        // testing exception
      }
  
  
      host.addDesiredConfig(c1.getClusterId(), true, "_test", config);
  
      Map<String, DesiredConfig> map = host.getDesiredConfigs(c1.getClusterId());
      Assert.assertTrue("Expect desired config to contain global", map.containsKey("global"));
  
      config = configFactory.createNew(c1, "global", "v2",
          new HashMap<String,String>() {{ put("c", "d"); }}, new HashMap<>());
      host.addDesiredConfig(c1.getClusterId(), true, "_test1", config);
  
      map = host.getDesiredConfigs(c1.getClusterId());
      Assert.assertTrue("Expect desired config to contain global", map.containsKey("global"));
      Assert.assertEquals("Expect version to be 'v2'", "v2", map.get("global").getTag());
  
      host.addDesiredConfig(c1.getClusterId(), false, "_test2", config);
      map = host.getDesiredConfigs(c1.getClusterId());
      Assert.assertEquals("Expect no mapping configs", 0, map.size());
  
    }
  
    @Test
    public void testHostMaintenance() throws Exception {
      AmbariMetaInfo metaInfo = injector.getInstance(AmbariMetaInfo.class);
  
      StackId stackId = new StackId("HDP-0.1");
      clusters.addCluster("c1", stackId);
      Cluster c1 = clusters.getCluster("c1");
      Assert.assertEquals("c1", c1.getClusterName());
      clusters.addHost("h1");
      Host host = clusters.getHost("h1");
      host.setIPv4("ipv4");
      host.setIPv6("ipv6");
  
      Map<String, String> hostAttributes = new HashMap<>();
      hostAttributes.put("os_family", "redhat");
      hostAttributes.put("os_release_version", "6.3");
      host.setHostAttributes(hostAttributes);
  
      helper.getOrCreateRepositoryVersion(stackId, stackId.getStackVersion());
      c1.setDesiredStackVersion(stackId);
      clusters.mapHostToCluster("h1", "c1");
  
      HostEntity entity = hostDAO.findByName("h1");
      HostStateEntity stateEntity = entity.getHostStateEntity();
      Assert.assertNull(stateEntity.getMaintenanceState());
      Assert.assertEquals(MaintenanceState.OFF, host.getMaintenanceState(c1.getClusterId()));
  
      host.setMaintenanceState(c1.getClusterId(), MaintenanceState.ON);
  
      entity = hostDAO.findByName("h1");
      stateEntity = entity.getHostStateEntity();
      Assert.assertNotNull(stateEntity.getMaintenanceState());
      Assert.assertEquals(MaintenanceState.ON, host.getMaintenanceState(c1.getClusterId()));
    }
  
    @Test
    public void testHostPersist() throws Exception {
      clusters.addHost("foo");
      Host host = clusters.getHost("foo");
  
      String rackInfo = "rackInfo";
      long lastRegistrationTime = System.currentTimeMillis();
  
      host.setRackInfo(rackInfo);
      host.setLastRegistrationTime(lastRegistrationTime);
  
      Assert.assertEquals(rackInfo, host.getRackInfo());
      Assert.assertEquals(lastRegistrationTime, host.getLastRegistrationTime());
    }
  }
  