heron/healthmgr/tests/java/com/twitter/heron/healthmgr/HealthManagerTest.java

// Copyright 2016 Twitter. All rights reserved.
//
// Licensed 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 com.twitter.heron.healthmgr;

import java.util.Arrays;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

import javax.inject.Inject;

import com.google.inject.AbstractModule;
import com.microsoft.dhalion.api.IHealthPolicy;
import com.microsoft.dhalion.api.MetricsProvider;
import com.microsoft.dhalion.policy.HealthPolicyImpl;

import org.junit.Test;

import com.twitter.heron.healthmgr.HealthPolicyConfigReader.PolicyConfigKey;
import com.twitter.heron.healthmgr.sensors.TrackerMetricsProvider;
import com.twitter.heron.proto.scheduler.Scheduler.SchedulerLocation;
import com.twitter.heron.scheduler.client.ISchedulerClient;
import com.twitter.heron.spi.common.Config;
import com.twitter.heron.spi.common.Key;
import com.twitter.heron.spi.statemgr.SchedulerStateManagerAdaptor;

import static org.junit.Assert.assertEquals;
import static org.junit.Assert.assertNotNull;
import static org.mockito.Mockito.anyString;
import static org.mockito.Mockito.doReturn;
import static org.mockito.Mockito.mock;
import static org.mockito.Mockito.spy;
import static org.mockito.Mockito.when;

public class HealthManagerTest {
  @Test
  public void testInitialize() throws Exception {
    String topologyName = "testTopology";
    Config config = Config.newBuilder()
        .put(Key.SCHEDULER_IS_SERVICE, true)
        .put(Key.TOPOLOGY_NAME, topologyName)
        .put(Key.ENVIRON, "environ")
        .put(Key.CLUSTER, "cluster")
        .build();

    SchedulerStateManagerAdaptor adaptor = mock(SchedulerStateManagerAdaptor.class);

    SchedulerLocation schedulerLocation = SchedulerLocation.newBuilder()
        .setTopologyName(topologyName)
        .setHttpEndpoint("http://127.0.0.1")
        .build();
    when(adaptor.getSchedulerLocation(anyString())).thenReturn(schedulerLocation);

    AbstractModule baseModule = HealthManager
        .buildMetricsProviderModule("127.0.0.1", TrackerMetricsProvider.class.getName());

    HealthManager healthManager = new HealthManager(config, baseModule);

    Map<String, Object> policy = new HashMap<>();
    policy.put(PolicyConfigKey.HEALTH_POLICY_CLASS.key(), TestPolicy.class.getName());
    policy.put("test-config", "test-value");

    String[] policyIds = {"policy"};

    HealthPolicyConfigReader policyConfigProvider = mock(HealthPolicyConfigReader.class);
    when(policyConfigProvider.getPolicyIds()).thenReturn(Arrays.asList(policyIds));
    when(policyConfigProvider.getPolicyConfig("policy")).thenReturn(policy);

    HealthManager spyHealthMgr = spy(healthManager);
    doReturn(adaptor).when(spyHealthMgr).createStateMgrAdaptor();
    doReturn(policyConfigProvider).when(spyHealthMgr).createPolicyConfigReader();

    spyHealthMgr.initialize();

    List<IHealthPolicy> healthPolicies = spyHealthMgr.getHealthPolicies();
    assertEquals(1, healthPolicies.size());
    TestPolicy healthPolicy = (TestPolicy) healthPolicies.iterator().next();
    assertNotNull(healthPolicy.schedulerClient);
    assertEquals(healthPolicy.stateMgrAdaptor, adaptor);
    assertNotNull(healthPolicy.metricsProvider);
    assertEquals(healthPolicy.config.getConfig("test-config"), "test-value");
  }

  static class TestPolicy extends HealthPolicyImpl {
    private HealthPolicyConfig config;
    private final ISchedulerClient schedulerClient;
    private final SchedulerStateManagerAdaptor stateMgrAdaptor;
    private final MetricsProvider metricsProvider;

    @Inject
    public TestPolicy(HealthPolicyConfig config,
                      ISchedulerClient schedulerClient,
                      SchedulerStateManagerAdaptor stateMgrAdaptor,
                      MetricsProvider metricsProvider) {
      this.config = config;
      this.schedulerClient = schedulerClient;
      this.stateMgrAdaptor = stateMgrAdaptor;
      this.metricsProvider = metricsProvider;
    }
  }
}