org.apache.helix.ConfigAccessor Java Examples

@Test
public void testAddCustomizedStateConfig() {
  String className = TestHelper.getTestClassName();
  String methodName = TestHelper.getTestMethodName();
  String clusterName = className + "_" + methodName;

  HelixAdmin admin = new ZKHelixAdmin(ZK_ADDR);
  admin.addCluster(clusterName, true);
  CustomizedStateConfig.Builder builder =
      new CustomizedStateConfig.Builder();
  builder.addAggregationEnabledType("mockType1");
  CustomizedStateConfig customizedStateConfig = builder.build();

  admin.addCustomizedStateConfig(clusterName, customizedStateConfig);

  // Read CustomizedStateConfig from Zookeeper and check the content
  ConfigAccessor _configAccessor = new ConfigAccessor(ZK_ADDR);
  CustomizedStateConfig configFromZk =
      _configAccessor.getCustomizedStateConfig(clusterName);
  List<String> listTypesFromZk = configFromZk.getAggregationEnabledTypes();
  Assert.assertEquals(listTypesFromZk.get(0), "mockType1");
}
 

@Test
public void testNoControllerLeader() {
  HelixManager helixManager = mock(HelixManager.class);
  HelixDataAccessor helixDataAccessor = mock(HelixDataAccessor.class);

  // Mock that there is no helix leader.
  when(helixManager.getHelixDataAccessor()).thenReturn(helixDataAccessor);
  PropertyKey.Builder keyBuilder = mock(PropertyKey.Builder.class);
  when(helixDataAccessor.keyBuilder()).thenReturn(keyBuilder);
  PropertyKey controllerLeader = mock(PropertyKey.class);
  when(keyBuilder.controllerLeader()).thenReturn(controllerLeader);
  when(helixDataAccessor.getProperty(controllerLeader)).thenReturn(null);

  // Lead controller resource disabled.
  ConfigAccessor configAccessor = mock(ConfigAccessor.class);
  ResourceConfig resourceConfig = mock(ResourceConfig.class);
  when(helixManager.getConfigAccessor()).thenReturn(configAccessor);
  when(configAccessor.getResourceConfig(any(), any())).thenReturn(resourceConfig);
  when(resourceConfig.getSimpleConfig(anyString())).thenReturn("false");

  // Create Controller Leader Locator
  FakeControllerLeaderLocator.create(helixManager);
  ControllerLeaderLocator controllerLeaderLocator = FakeControllerLeaderLocator.getInstance();

  Assert.assertNull(controllerLeaderLocator.getControllerLeader(testTable));
}
 

@Test(dependsOnMethods = "testAddClusterWithValidCustomizedCloudConfig")
public void testAddClusterWithCloudConfigDisabledCloud() throws Exception {
  String className = TestHelper.getTestClassName();
  String methodName = TestHelper.getTestMethodName();
  String clusterName = className + "_" + methodName;

  ZNRecord record = new ZNRecord("testZnode");
  record.setBooleanField(CloudConfig.CloudConfigProperty.CLOUD_ENABLED.name(), false);
  record.setSimpleField(CloudConfig.CloudConfigProperty.CLOUD_ID.name(), "TestCloudID");
  record.setSimpleField(CloudConfig.CloudConfigProperty.CLOUD_PROVIDER.name(),
      CloudProvider.AZURE.name());

  Map<String, String> map = new HashMap<>();
  map.put("addCloudConfig", "true");

  put("clusters/" + clusterName, map,
      Entity.entity(OBJECT_MAPPER.writeValueAsString(record), MediaType.APPLICATION_JSON_TYPE),
      Response.Status.CREATED.getStatusCode());

  // Read CloudConfig from Zookeeper and check the content
  ConfigAccessor _configAccessor = new ConfigAccessor(ZK_ADDR);
  CloudConfig cloudConfigFromZk = _configAccessor.getCloudConfig(clusterName);
  Assert.assertFalse(cloudConfigFromZk.isCloudEnabled());
  Assert.assertEquals(cloudConfigFromZk.getCloudID(), "TestCloudID");
  Assert.assertEquals(cloudConfigFromZk.getCloudProvider(), CloudProvider.AZURE.name());
}
 

@GET
@Path("{clusterId}/customized-state-config")
public Response getCustomizedStateConfig(@PathParam("clusterId") String clusterId) {
  if (!doesClusterExist(clusterId)) {
    return notFound(String.format("Cluster %s does not exist", clusterId));
  }

  ConfigAccessor configAccessor = getConfigAccessor();
  CustomizedStateConfig customizedStateConfig =
      configAccessor.getCustomizedStateConfig(clusterId);

  if (customizedStateConfig != null) {
    return JSONRepresentation(customizedStateConfig.getRecord());
  }

  return notFound();
}
 

public void addCluster(String clusterName, boolean overwritePrevious, CloudConfig cloudConfig)
    throws HelixException {
  _admin.addCluster(clusterName, overwritePrevious);
  for (BuiltInStateModelDefinitions def : BuiltInStateModelDefinitions.values()) {
    addStateModelDef(clusterName, def.getStateModelDefinition().getId(),
        def.getStateModelDefinition(), overwritePrevious);
  }

  if (cloudConfig != null) {
    _admin.addCloudConfig(clusterName, cloudConfig);
    // If cloud is enabled and Cloud Provider is Azure, populated the Topology information in cluster config
    if (cloudConfig.isCloudEnabled()
        && cloudConfig.getCloudProvider().equals(CloudProvider.AZURE.name())) {
      ConfigAccessor configAccessor = new ConfigAccessor(_zkClient);
      ClusterConfig clusterConfig = new ClusterConfig(clusterName);
      clusterConfig.setTopology(AzureConstants.AZURE_TOPOLOGY);
      clusterConfig.setTopologyAwareEnabled(true);
      clusterConfig.setFaultZoneType(AzureConstants.AZURE_FAULT_ZONE_TYPE);
      configAccessor.updateClusterConfig(clusterName, clusterConfig);
    }
  }
}
 

@Transition(to = "ONLINE", from = "OFFLINE")
public void onBecomeOnlineFromOffline(Message message, NotificationContext context)
    throws Exception {
  LOG.debug(_workerId + " becomes ONLINE from OFFLINE for " + _partition);
  ConfigAccessor clusterConfig = context.getManager().getConfigAccessor();
  HelixManager manager = context.getManager();
  HelixConfigScope clusterScope =
      new HelixConfigScopeBuilder(ConfigScopeProperty.CLUSTER).forCluster(
          manager.getClusterName()).build();
  String json = clusterConfig.get(clusterScope, message.getResourceName());
  Dag.Node node = Dag.Node.fromJson(json);
  Set<String> parentIds = node.getParentIds();
  String resourceName = message.getResourceName();
  int numPartitions = node.getNumPartitions();
  Task task = _taskFactory.createTask(resourceName, parentIds, manager, _taskResultStore);
  manager.addExternalViewChangeListener(task);

  LOG.debug("Starting task for " + _partition + "...");
  int partitionNum = Integer.parseInt(_partition.split("_")[1]);
  task.execute(resourceName, numPartitions, partitionNum);
  LOG.debug("Task for " + _partition + " done");
}
 

private void verifyConfigAccessorMsdsEndpoint(
    RealmAwareZkClient.RealmAwareZkConnectionConfig connectionConfig) {
  System.out.println("Start " + TestHelper.getTestMethodName());

  ConfigAccessor firstConfigAccessor = new ConfigAccessor.Builder().build();
  ConfigAccessor secondConfigAccessor =
      new ConfigAccessor.Builder().setRealmAwareZkConnectionConfig(connectionConfig).build();

  try {
    verifyMsdsZkRealm(CLUSTER_ONE, true,
        () -> firstConfigAccessor.getClusterConfig(CLUSTER_ONE));
    verifyMsdsZkRealm(CLUSTER_FOUR, false,
        () -> firstConfigAccessor.getClusterConfig(CLUSTER_FOUR));

    verifyMsdsZkRealm(CLUSTER_FOUR, true,
        () -> secondConfigAccessor.getClusterConfig(CLUSTER_FOUR));
    verifyMsdsZkRealm(CLUSTER_ONE, false,
        () -> secondConfigAccessor.getClusterConfig(CLUSTER_ONE));
  } finally {
    firstConfigAccessor.close();
    secondConfigAccessor.close();
  }
}
 

@BeforeClass
public void beforeClass() throws Exception {
  _participants = new MockParticipantManager[_numNodes];
  _numDbs = 1;
  _numPartitions = 20;
  _numNodes = 2;
  _numReplicas = 2;
  _verifier =
      new BestPossibleExternalViewVerifier.Builder(CLUSTER_NAME).setZkAddr(ZK_ADDR).build();

  _gSetupTool.addCluster(CLUSTER_NAME, true);
  setupParticipants();
  setupDBs();

  registerParticipants(_participants, _numNodes, _startPort, 0, -3000000L);

  // start controller
  String controllerName = CONTROLLER_PREFIX + "_0";
  _controller = new ClusterControllerManager(ZK_ADDR, CLUSTER_NAME, controllerName);
  _controller.syncStart();

  createManagers();
  _configAccessor = new ConfigAccessor(_gZkClient);
}
 

@GET
@Path("{clusterId}/cloudconfig")
public Response getCloudConfig(@PathParam("clusterId") String clusterId) {

  RealmAwareZkClient zkClient = getRealmAwareZkClient();
  if (!ZKUtil.isClusterSetup(clusterId, zkClient)) {
    return notFound();
  }

  ConfigAccessor configAccessor = new ConfigAccessor(zkClient);
  CloudConfig cloudConfig = configAccessor.getCloudConfig(clusterId);

  if (cloudConfig != null) {
    return JSONRepresentation(cloudConfig.getRecord());
  }

  return notFound();
}
 

@Test(dependsOnMethods = "testActivateSuperCluster")
public void testAddClusterWithCloudConfig() throws Exception {
  String className = TestHelper.getTestClassName();
  String methodName = TestHelper.getTestMethodName();
  String clusterName = className + "_" + methodName;

  ZNRecord record = new ZNRecord("testZnode");
  record.setBooleanField(CloudConfig.CloudConfigProperty.CLOUD_ENABLED.name(), true);
  record.setSimpleField(CloudConfig.CloudConfigProperty.CLOUD_ID.name(), "TestCloudID");
  record.setSimpleField(CloudConfig.CloudConfigProperty.CLOUD_PROVIDER.name(),
      CloudProvider.AZURE.name());

  Map<String, String> map = new HashMap<>();
  map.put("addCloudConfig", "true");

  put("clusters/" + clusterName, map,
      Entity.entity(OBJECT_MAPPER.writeValueAsString(record), MediaType.APPLICATION_JSON_TYPE),
      Response.Status.CREATED.getStatusCode());

  // Read CloudConfig from Zookeeper and check the content
  ConfigAccessor _configAccessor = new ConfigAccessor(ZK_ADDR);
  CloudConfig cloudConfigFromZk = _configAccessor.getCloudConfig(clusterName);
  Assert.assertTrue(cloudConfigFromZk.isCloudEnabled());
  Assert.assertEquals(cloudConfigFromZk.getCloudID(), "TestCloudID");
  Assert.assertEquals(cloudConfigFromZk.getCloudProvider(), CloudProvider.AZURE.name());

  ClusterConfig clusterConfigFromZk = _configAccessor.getClusterConfig(clusterName);
  Assert.assertEquals(clusterConfigFromZk.getTopology(), AzureConstants.AZURE_TOPOLOGY);
  Assert.assertEquals(clusterConfigFromZk.getFaultZoneType(), AzureConstants.AZURE_FAULT_ZONE_TYPE);
  Assert.assertTrue(clusterConfigFromZk.isTopologyAwareEnabled());
}
 

@BeforeClass
public void beforeClass() throws Exception {
  System.out.println("START " + CLASS_NAME + " at " + new Date(System.currentTimeMillis()));

  _gSetupTool.addCluster(CLUSTER_NAME, true);

  for (int i = 0; i < NUM_NODE; i++) {
    String storageNodeName = PARTICIPANT_PREFIX + "_" + (START_PORT + i);
    _gSetupTool.addInstanceToCluster(CLUSTER_NAME, storageNodeName);

    // start dummy participants
    MockParticipantManager participant =
        new MockParticipantManager(ZK_ADDR, CLUSTER_NAME, storageNodeName);
    participant.syncStart();
    _participants.add(participant);
  }

  // start controller
  String controllerName = CONTROLLER_PREFIX + "_0";
  _controller = new ClusterControllerManager(ZK_ADDR, CLUSTER_NAME, controllerName);
  _controller.syncStart();

  _clusterVerifier =
      new BestPossibleExternalViewVerifier.Builder(CLUSTER_NAME).setZkAddr(ZK_ADDR).build();

  enablePersistBestPossibleAssignment(_gZkClient, CLUSTER_NAME, true);

  _configAccessor = new ConfigAccessor(_gZkClient);
}
 

@Test
public void testQuotaConstraintSkipWorkflowAssignment() throws Exception {
  ClusterEvent event = new ClusterEvent(ClusterEventType.Unknown);
  WorkflowControllerDataProvider cache = new WorkflowControllerDataProvider(CLUSTER_NAME);
  JobConfig.Builder job = new JobConfig.Builder();

  job.setJobCommandConfigMap(Collections.singletonMap(MockTask.JOB_DELAY, "100000"));
  TaskDriver driver = new TaskDriver(_manager);
  for (int i = 0; i < 10; i++) {
    Workflow.Builder workflow = new Workflow.Builder("Workflow" + i);
    job.setWorkflow("Workflow" + i);
    TaskConfig taskConfig =
        new TaskConfig(MockTask.TASK_COMMAND, new HashMap<String, String>(), null, null);
    job.addTaskConfigMap(Collections.singletonMap(taskConfig.getId(), taskConfig));
    job.setJobId(TaskUtil.getNamespacedJobName("Workflow" + i, "JOB"));
    workflow.addJob("JOB", job);
    driver.start(workflow.build());
  }
  ConfigAccessor accessor = new ConfigAccessor(_gZkClient);
  ClusterConfig clusterConfig = accessor.getClusterConfig(CLUSTER_NAME);
  clusterConfig.setTaskQuotaRatio(AssignableInstance.DEFAULT_QUOTA_TYPE, 3);
  clusterConfig.setTaskQuotaRatio("OtherType", 37);
  accessor.setClusterConfig(CLUSTER_NAME, clusterConfig);
  cache.refresh(_manager.getHelixDataAccessor());
  event.addAttribute(AttributeName.ControllerDataProvider.name(), cache);
  event.addAttribute(AttributeName.helixmanager.name(), _manager);
  runStage(event, new ResourceComputationStage());
  runStage(event, new CurrentStateComputationStage());
  runStage(event, new TaskSchedulingStage());
  Assert.assertTrue(!cache.getAssignableInstanceManager()
      .hasGlobalCapacity(AssignableInstance.DEFAULT_QUOTA_TYPE));
  BestPossibleStateOutput bestPossibleStateOutput =
      event.getAttribute(AttributeName.BEST_POSSIBLE_STATE.name());
  Assert.assertTrue(bestPossibleStateOutput.getStateMap().size() == 3);
}
 

@Test(dependsOnMethods = "testCustomizedStateConfig")
public void testCustomizedStateConfigBuilder() {
  String className = getShortClassName();
  String clusterName = "CLUSTER_" + className;
  TestHelper.setupEmptyCluster(_gZkClient, clusterName);
  CustomizedStateConfig.Builder builder =
      new CustomizedStateConfig.Builder();
  builder.addAggregationEnabledType("mockType1");
  builder.addAggregationEnabledType("mockType2");

  // Check builder getter methods
  List<String> aggregationEnabledTypes = builder.getAggregationEnabledTypes();
  Assert.assertEquals(aggregationEnabledTypes.size(), 2);
  Assert.assertEquals(aggregationEnabledTypes.get(0), "mockType1");
  Assert.assertEquals(aggregationEnabledTypes.get(1), "mockType2");

  CustomizedStateConfig customizedStateConfig = builder.build();

  ZKHelixDataAccessor accessor =
      new ZKHelixDataAccessor(clusterName, new ZkBaseDataAccessor(ZK_ADDR));
  Builder keyBuilder = accessor.keyBuilder();
  accessor.setProperty(keyBuilder.customizedStateConfig(),
      customizedStateConfig);

  // Read CustomizedStateConfig from Zookeeper and check the content
  ConfigAccessor _configAccessor = new ConfigAccessor(ZK_ADDR);
  CustomizedStateConfig customizedStateConfigFromZk =
      _configAccessor.getCustomizedStateConfig(clusterName);
  List<String> aggregationEnabledTypesFromZk =
      customizedStateConfigFromZk.getAggregationEnabledTypes();
  Assert.assertEquals(aggregationEnabledTypesFromZk.get(0), "mockType1");
  Assert.assertEquals(aggregationEnabledTypesFromZk.get(1), "mockType2");
}
 

@Test(expectedExceptions = HelixException.class)
public void testCloudConfigNonExistentCluster() {
  String className = getShortClassName();
  String clusterName = "CLUSTER_" + className;
  // Read CloudConfig from Zookeeper and get exception since cluster in not setup yet
  ConfigAccessor _configAccessor = new ConfigAccessor(_gZkClient);
  CloudConfig cloudConfigFromZk = _configAccessor.getCloudConfig(clusterName);
}
 

@Test(dependsOnMethods = "testCloudConfigNonExistentCluster")
public void testCloudConfigNull() {
  String className = getShortClassName();
  String clusterName = "CLUSTER_" + className;
  TestHelper.setupEmptyCluster(_gZkClient, clusterName);
  // Read CloudConfig from Zookeeper
  ConfigAccessor _configAccessor = new ConfigAccessor(_gZkClient);
  CloudConfig cloudConfigFromZk = _configAccessor.getCloudConfig(clusterName);
  // since CloudConfig is not written to ZooKeeper, the output should be null
  Assert.assertNull(cloudConfigFromZk);
}
 

@Test(dependsOnMethods = "testAddCustomizedConfig")
public void testDeleteCustomizedConfig() throws IOException {
  System.out.println("Start test :" + TestHelper.getTestMethodName());
  _gSetupTool.addCluster("TestClusterCustomized", true);
  String urlBase = "clusters/TestClusterCustomized/customized-state-config/";
  ZNRecord record = new ZNRecord("TestCustomizedStateConfig");
  List<String> testList = new ArrayList<String>();
  testList.add("mockType1");
  record.setListField(
      CustomizedStateConfig.CustomizedStateProperty.AGGREGATION_ENABLED_TYPES
          .name(),
      testList);

  put(urlBase, null,
      Entity.entity(OBJECT_MAPPER.writeValueAsString(record), MediaType.APPLICATION_JSON_TYPE),
      Response.Status.OK.getStatusCode());

  // Read CustomizedStateConfig from Zookeeper and make sure it exists
  ConfigAccessor _configAccessor = new ConfigAccessor(ZK_ADDR);
  CustomizedStateConfig customizedConfigFromZk = _configAccessor.getCustomizedStateConfig("TestClusterCustomized");
  Assert.assertNotNull(customizedConfigFromZk);

  delete(urlBase, Response.Status.OK.getStatusCode());

  customizedConfigFromZk = _configAccessor.getCustomizedStateConfig("TestClusterCustomized");
  Assert.assertNull(customizedConfigFromZk);

  System.out.println("End test :" + TestHelper.getTestMethodName());
}
 

/**
 * Populate info on ZooKeeper server and start {@link HelixControllerManager}.
 * @param zkConnectString zk connect string to zk server.
 * @param vcrClusterName the vcr cluster name.
 * @param clusterMap the {@link ClusterMap} to use.
 * @return the created {@link HelixControllerManager}.
 */
public static HelixControllerManager populateZkInfoAndStartController(String zkConnectString, String vcrClusterName,
    ClusterMap clusterMap) {
  HelixZkClient zkClient = DedicatedZkClientFactory.getInstance()
      .buildZkClient(new HelixZkClient.ZkConnectionConfig(zkConnectString), new HelixZkClient.ZkClientConfig());
  try {
    zkClient.setZkSerializer(new ZNRecordSerializer());
    ClusterSetup clusterSetup = new ClusterSetup(zkClient);
    clusterSetup.addCluster(vcrClusterName, true);
    HelixAdmin admin = new ZKHelixAdmin(zkClient);
    // set ALLOW_PARTICIPANT_AUTO_JOIN
    HelixConfigScope configScope = new HelixConfigScopeBuilder(HelixConfigScope.ConfigScopeProperty.CLUSTER).
        forCluster(vcrClusterName).build();
    Map<String, String> helixClusterProperties = new HashMap<>();
    helixClusterProperties.put(ZKHelixManager.ALLOW_PARTICIPANT_AUTO_JOIN, String.valueOf(true));
    admin.setConfig(configScope, helixClusterProperties);
    // set PersistBestPossibleAssignment
    ConfigAccessor configAccessor = new ConfigAccessor(zkClient);
    ClusterConfig clusterConfig = configAccessor.getClusterConfig(vcrClusterName);
    clusterConfig.setPersistBestPossibleAssignment(true);
    configAccessor.setClusterConfig(vcrClusterName, clusterConfig);

    FullAutoModeISBuilder builder = new FullAutoModeISBuilder(helixResource);
    builder.setStateModel(LeaderStandbySMD.name);
    for (PartitionId partitionId : clusterMap.getAllPartitionIds(null)) {
      builder.add(partitionId.toPathString());
    }
    builder.setRebalanceStrategy(CrushEdRebalanceStrategy.class.getName());
    IdealState idealState = builder.build();
    admin.addResource(vcrClusterName, helixResource, idealState);
    admin.rebalance(vcrClusterName, helixResource, 3, "", "");
    HelixControllerManager helixControllerManager = new HelixControllerManager(zkConnectString, vcrClusterName);
    helixControllerManager.syncStart();
    return helixControllerManager;
  } finally {
    zkClient.close();
  }
}
 

@Test(dependsOnMethods = "testCustomizedStateConfigNull")
public void testCustomizedStateConfig() {
  String className = getShortClassName();
  String clusterName = "CLUSTER_" + className;
  TestHelper.setupEmptyCluster(_gZkClient, clusterName);

  // Create dummy CustomizedStateConfig object
  CustomizedStateConfig.Builder customizedStateConfigBuilder =
      new CustomizedStateConfig.Builder();
  List<String> aggregationEnabledTypes = new ArrayList<String>();
  aggregationEnabledTypes.add("mockType1");
  aggregationEnabledTypes.add("mockType2");
  customizedStateConfigBuilder.setAggregationEnabledTypes(aggregationEnabledTypes);
  CustomizedStateConfig customizedStateConfig =
      customizedStateConfigBuilder.build();

  // Write the CustomizedStateConfig to Zookeeper
  ZKHelixDataAccessor accessor =
      new ZKHelixDataAccessor(clusterName, new ZkBaseDataAccessor(ZK_ADDR));
  Builder keyBuilder = accessor.keyBuilder();
  accessor.setProperty(keyBuilder.customizedStateConfig(),
      customizedStateConfig);

  // Read CustomizedStateConfig from Zookeeper and check the content
  ConfigAccessor _configAccessor = new ConfigAccessor(ZK_ADDR);
  CustomizedStateConfig customizedStateConfigFromZk =
      _configAccessor.getCustomizedStateConfig(clusterName);
  Assert.assertEquals(customizedStateConfigFromZk.getAggregationEnabledTypes().size(),
      2);
  Assert.assertEquals(aggregationEnabledTypes.get(0), "mockType1");
  Assert.assertEquals(aggregationEnabledTypes.get(1), "mockType2");
}
 

@Test
public void testRemoveCustomizedStateConfig() throws Exception {
  String className = TestHelper.getTestClassName();
  String methodName = TestHelper.getTestMethodName();
  String clusterName = className + "_" + methodName;

  HelixAdmin admin = new ZKHelixAdmin(ZK_ADDR);
  admin.addCluster(clusterName, true);
  CustomizedStateConfig.Builder builder =
      new CustomizedStateConfig.Builder();
  builder.addAggregationEnabledType("mockType1");
  CustomizedStateConfig customizedStateConfig = builder.build();

  admin.addCustomizedStateConfig(clusterName, customizedStateConfig);

  // Read CustomizedStateConfig from Zookeeper and check the content
  ConfigAccessor _configAccessor = new ConfigAccessor(ZK_ADDR);
  CustomizedStateConfig configFromZk =
      _configAccessor.getCustomizedStateConfig(clusterName);
  List<String> listTypesFromZk = configFromZk.getAggregationEnabledTypes();
  Assert.assertEquals(listTypesFromZk.get(0), "mockType1");

  // Remove CustomizedStateConfig Config and make sure it has been removed from
  // Zookeeper
  admin.removeCustomizedStateConfig(clusterName);
  configFromZk = _configAccessor.getCustomizedStateConfig(clusterName);
  Assert.assertNull(configFromZk);
}
 

@Test
public void testRemoveCloudConfig() throws Exception {
  String className = TestHelper.getTestClassName();
  String methodName = TestHelper.getTestMethodName();
  String clusterName = className + "_" + methodName;

  HelixAdmin admin = new ZKHelixAdmin(_gZkClient);
  admin.addCluster(clusterName, true);

  CloudConfig.Builder builder = new CloudConfig.Builder();
  builder.setCloudEnabled(true);
  builder.setCloudID("TestID");
  builder.addCloudInfoSource("TestURL");
  builder.setCloudProvider(CloudProvider.CUSTOMIZED);
  builder.setCloudInfoProcessorName("TestProcessor");
  CloudConfig cloudConfig = builder.build();

  admin.addCloudConfig(clusterName, cloudConfig);

  // Read CloudConfig from Zookeeper and check the content
  ConfigAccessor _configAccessor = new ConfigAccessor(_gZkClient);
  CloudConfig cloudConfigFromZk = _configAccessor.getCloudConfig(clusterName);
  Assert.assertTrue(cloudConfigFromZk.isCloudEnabled());
  Assert.assertEquals(cloudConfigFromZk.getCloudID(), "TestID");
  Assert.assertEquals(cloudConfigFromZk.getCloudProvider(), CloudProvider.CUSTOMIZED.name());
  List<String> listUrlFromZk = cloudConfigFromZk.getCloudInfoSources();
  Assert.assertEquals(listUrlFromZk.get(0), "TestURL");
  Assert.assertEquals(cloudConfigFromZk.getCloudInfoProcessorName(), "TestProcessor");

  // Remove Cloud Config and make sure it has been removed from Zookeeper
  admin.removeCloudConfig(clusterName);
  cloudConfigFromZk = _configAccessor.getCloudConfig(clusterName);
  Assert.assertNull(cloudConfigFromZk);
}
 

@Test
public void testAddCloudConfig() {
  String className = TestHelper.getTestClassName();
  String methodName = TestHelper.getTestMethodName();
  String clusterName = className + "_" + methodName;

  HelixAdmin admin = new ZKHelixAdmin(_gZkClient);
  admin.addCluster(clusterName, true);

  CloudConfig.Builder builder = new CloudConfig.Builder();
  builder.setCloudEnabled(true);
  builder.setCloudID("TestID");
  builder.addCloudInfoSource("TestURL");
  builder.setCloudProvider(CloudProvider.CUSTOMIZED);
  builder.setCloudInfoProcessorName("TestProcessor");
  CloudConfig cloudConfig = builder.build();

  admin.addCloudConfig(clusterName, cloudConfig);

  // Read CloudConfig from Zookeeper and check the content
  ConfigAccessor _configAccessor = new ConfigAccessor(_gZkClient);
  CloudConfig cloudConfigFromZk = _configAccessor.getCloudConfig(clusterName);
  Assert.assertTrue(cloudConfigFromZk.isCloudEnabled());
  Assert.assertEquals(cloudConfigFromZk.getCloudID(), "TestID");
  Assert.assertEquals(cloudConfigFromZk.getCloudProvider(), CloudProvider.CUSTOMIZED.name());
  List<String> listUrlFromZk = cloudConfigFromZk.getCloudInfoSources();
  Assert.assertEquals(listUrlFromZk.get(0), "TestURL");
  Assert.assertEquals(cloudConfigFromZk.getCloudInfoProcessorName(), "TestProcessor");
}
 

@VisibleForTesting
InstanceServiceImpl(HelixDataAccessorWrapper dataAccessor, ConfigAccessor configAccessor,
    CustomRestClient customRestClient) {
  _dataAccessor = dataAccessor;
  _configAccessor = configAccessor;
  _customRestClient = customRestClient;
}
 

private void setResourceThreadPoolSize(String resourceName, int threadPoolSize) {
  HelixManager manager = _participants[0];
  ConfigAccessor accessor = manager.getConfigAccessor();
  HelixConfigScope scope =
      new HelixConfigScopeBuilder(HelixConfigScope.ConfigScopeProperty.RESOURCE)
          .forCluster(manager.getClusterName()).forResource(resourceName).build();
  accessor.set(scope, HelixTaskExecutor.MAX_THREADS, "" + threadPoolSize);
}
 

@BeforeClass
public void beforeClass() throws Exception {
  System.out.println("START " + CLASS_NAME + " at " + new Date(System.currentTimeMillis()));

  _gSetupTool.addCluster(CLUSTER_NAME, true);

  for (int i = 0; i < NUM_NODE; i++) {
    String storageNodeName = PARTICIPANT_PREFIX + "_" + (START_PORT + i);
    MockParticipantManager participant = createAndStartParticipant(storageNodeName);
    _participants.add(participant);
  }

  // start controller
  String controllerName = CONTROLLER_PREFIX + "_0";
  _controller = new ClusterControllerManager(ZK_ADDR, CLUSTER_NAME, controllerName);
  _controller.syncStart();

  _clusterVerifier =
      new BestPossibleExternalViewVerifier.Builder(CLUSTER_NAME).setZkAddr(ZK_ADDR).build();

  enablePersistIntermediateAssignment(_gZkClient, CLUSTER_NAME, true);

  _manager = HelixManagerFactory
      .getZKHelixManager(CLUSTER_NAME, "admin", InstanceType.ADMINISTRATOR, ZK_ADDR);
  _manager.connect();
  _configAccessor = new ConfigAccessor(_gZkClient);
}
 

@Test
public void testControllerLeaderExists() {
  HelixManager helixManager = mock(HelixManager.class);
  HelixDataAccessor helixDataAccessor = mock(HelixDataAccessor.class);
  HelixAdmin helixAdmin = mock(HelixAdmin.class);
  final String leaderHost = "host";
  final int leaderPort = 12345;

  // Lead controller resource disabled.
  ConfigAccessor configAccessor = mock(ConfigAccessor.class);
  ResourceConfig resourceConfig = mock(ResourceConfig.class);
  when(helixManager.getConfigAccessor()).thenReturn(configAccessor);
  when(configAccessor.getResourceConfig(any(), anyString())).thenReturn(resourceConfig);
  when(resourceConfig.getSimpleConfig(anyString())).thenReturn("false");

  // Mocks the helix leader
  when(helixManager.getHelixDataAccessor()).thenReturn(helixDataAccessor);
  PropertyKey.Builder keyBuilder = mock(PropertyKey.Builder.class);
  when(helixDataAccessor.keyBuilder()).thenReturn(keyBuilder);
  PropertyKey controllerLeader = mock(PropertyKey.class);
  when(keyBuilder.controllerLeader()).thenReturn(controllerLeader);
  LiveInstance liveInstance = mock(LiveInstance.class);
  when(helixDataAccessor.getProperty(controllerLeader)).thenReturn(liveInstance);
  when(liveInstance.getInstanceName()).thenReturn(leaderHost + "_" + leaderPort);

  when(helixManager.getClusterName()).thenReturn("myCluster");
  when(helixManager.getClusterManagmentTool()).thenReturn(helixAdmin);
  when(helixAdmin.getResourceExternalView(anyString(), anyString())).thenReturn(null);

  // Create Controller Leader Locator
  FakeControllerLeaderLocator.create(helixManager);
  ControllerLeaderLocator controllerLeaderLocator = FakeControllerLeaderLocator.getInstance();

  Pair<String, Integer> expectedLeaderLocation = new Pair<>(leaderHost, leaderPort);
  Assert.assertEquals(controllerLeaderLocator.getControllerLeader(testTable).getFirst(),
      expectedLeaderLocation.getFirst());
  Assert.assertEquals(controllerLeaderLocator.getControllerLeader(testTable).getSecond(),
      expectedLeaderLocation.getSecond());
}
 

@Test(dependsOnMethods = "testAddClusterWithInvalidCloudConfig")
public void testAddClusterWithValidCloudConfig() throws Exception {
  String className = TestHelper.getTestClassName();
  String methodName = TestHelper.getTestMethodName();
  String clusterName = className + "_" + methodName;

  CloudConfig.Builder cloudConfigInitBuilder = new CloudConfig.Builder();
  cloudConfigInitBuilder.setCloudEnabled(true);
  cloudConfigInitBuilder.setCloudID("TestID");
  List<String> sourceList = new ArrayList<String>();
  sourceList.add("TestURL");
  cloudConfigInitBuilder.setCloudInfoSources(sourceList);
  cloudConfigInitBuilder.setCloudInfoProcessorName("TestProcessorName");
  cloudConfigInitBuilder.setCloudProvider(CloudProvider.CUSTOMIZED);

  CloudConfig cloudConfigInit = cloudConfigInitBuilder.build();

  _clusterSetup.addCluster(clusterName, false, cloudConfigInit);

  // Read CloudConfig from Zookeeper and check the content
  ConfigAccessor _configAccessor = new ConfigAccessor(_gZkClient);
  CloudConfig cloudConfigFromZk = _configAccessor.getCloudConfig(clusterName);
  Assert.assertTrue(cloudConfigFromZk.isCloudEnabled());
  Assert.assertEquals(cloudConfigFromZk.getCloudID(), "TestID");
  List<String> listUrlFromZk = cloudConfigFromZk.getCloudInfoSources();
  Assert.assertEquals(listUrlFromZk.get(0), "TestURL");
  Assert.assertEquals(cloudConfigFromZk.getCloudInfoProcessorName(), "TestProcessorName");
  Assert.assertEquals(cloudConfigFromZk.getCloudProvider(), CloudProvider.CUSTOMIZED.name());
}
 

@Test(dependsOnMethods = "testAddClusterWithValidCloudConfig")
public void testAddClusterAzureProvider() throws Exception {
  String className = TestHelper.getTestClassName();
  String methodName = TestHelper.getTestMethodName();
  String clusterName = className + "_" + methodName;

  CloudConfig.Builder cloudConfigInitBuilder = new CloudConfig.Builder();
  cloudConfigInitBuilder.setCloudEnabled(true);
  cloudConfigInitBuilder.setCloudID("TestID");
  cloudConfigInitBuilder.setCloudProvider(CloudProvider.AZURE);

  CloudConfig cloudConfigInit = cloudConfigInitBuilder.build();

  _clusterSetup.addCluster(clusterName, false, cloudConfigInit);

  // Read CloudConfig from Zookeeper and check the content
  ConfigAccessor _configAccessor = new ConfigAccessor(ZK_ADDR);
  CloudConfig cloudConfigFromZk = _configAccessor.getCloudConfig(clusterName);
  Assert.assertTrue(cloudConfigFromZk.isCloudEnabled());
  Assert.assertEquals(cloudConfigFromZk.getCloudID(), "TestID");
  List<String> listUrlFromZk = cloudConfigFromZk.getCloudInfoSources();

  // Since it is Azure, topology information should have been populated.
  ClusterConfig clusterConfig = _configAccessor.getClusterConfig(clusterName);
  Assert.assertEquals(clusterConfig.getTopology(), AzureConstants.AZURE_TOPOLOGY);
  Assert.assertEquals(clusterConfig.getFaultZoneType(), AzureConstants.AZURE_FAULT_ZONE_TYPE);
  Assert.assertTrue(clusterConfig.isTopologyAwareEnabled());

  // Since provider is not customized, CloudInfoSources and CloudInfoProcessorName will be null.
  Assert.assertNull(listUrlFromZk);
  Assert.assertNull(cloudConfigFromZk.getCloudInfoProcessorName());
  Assert.assertEquals(cloudConfigFromZk.getCloudProvider(), CloudProvider.AZURE.name());
}
 

@BeforeSuite
public void beforeSuite()
    throws Exception {
  if (!_init) {
    setupZooKeepers();

    // TODO: use logging.properties file to config java.util.logging.Logger levels
    java.util.logging.Logger topJavaLogger = java.util.logging.Logger.getLogger("");
    topJavaLogger.setLevel(Level.WARNING);

    HelixZkClient.ZkClientConfig clientConfig = new HelixZkClient.ZkClientConfig();

    clientConfig.setZkSerializer(new ZNRecordSerializer());
    _gZkClient = DedicatedZkClientFactory.getInstance()
        .buildZkClient(new HelixZkClient.ZkConnectionConfig(ZK_ADDR), clientConfig);

    clientConfig.setZkSerializer(new ZNRecordSerializer());
    _gZkClientTestNS = DedicatedZkClientFactory.getInstance()
        .buildZkClient(new HelixZkClient.ZkConnectionConfig(_zkAddrTestNS), clientConfig);

    _gSetupTool = new ClusterSetup(_gZkClient);
    _configAccessor = new ConfigAccessor(_gZkClient);
    _baseAccessor = new ZkBaseDataAccessor<>(_gZkClient);
    _baseAccessorTestNS = new ZkBaseDataAccessor<>(_gZkClientTestNS);

    // wait for the web service to start
    Thread.sleep(100);
    setupHelixResources();
    _init = true;
  }
}
 

@Test
public void testDisableCluster() throws Exception {
  ConfigAccessor configAccessor = new ConfigAccessor(_gZkClient);
  ClusterConfig clusterConfig = configAccessor.getClusterConfig(CLUSTER_NAME);
  clusterConfig.setMaxPartitionsPerInstance(10);
  configAccessor.setClusterConfig(CLUSTER_NAME, clusterConfig);

  int i = 0;
  for (String stateModel : TestStateModels) {
    String db = "Test-DB-" + i++;
    int replica = 3;
    createResourceWithDelayedRebalance(CLUSTER_NAME, db, stateModel, _PARTITIONS, replica,
        replica, -1);
    _testDBs.add(db);
  }
  Thread.sleep(100L);
  Assert.assertTrue(_clusterVerifier.verifyByPolling());

  MaintenanceSignal maintenanceSignal =
      _dataAccessor.getProperty(_dataAccessor.keyBuilder().maintenance());
  Assert.assertNull(maintenanceSignal);

  for (i = 2; i < NUM_NODE; i++) {
    _participants.get(i).syncStop();
  }

  Thread.sleep(1000L);
  maintenanceSignal = _dataAccessor.getProperty(_dataAccessor.keyBuilder().maintenance());
  Assert.assertNotNull(maintenanceSignal);
  Assert.assertNotNull(maintenanceSignal.getReason());
}
 

@Test
public void testUpdateCustomizedStateConfig() throws Exception {
  String className = TestHelper.getTestClassName();
  String methodName = TestHelper.getTestMethodName();
  String clusterName = className + "_" + methodName;

  HelixAdmin admin = new ZKHelixAdmin(ZK_ADDR);
  admin.addCluster(clusterName, true);
  CustomizedStateConfig.Builder builder =
      new CustomizedStateConfig.Builder();
  builder.addAggregationEnabledType("mockType1");
  CustomizedStateConfig customizedStateConfig = builder.build();

  admin.addCustomizedStateConfig(clusterName, customizedStateConfig);

  // Read CustomizedStateConfig from Zookeeper and check the content
  ConfigAccessor _configAccessor = new ConfigAccessor(ZK_ADDR);
  CustomizedStateConfig configFromZk =
      _configAccessor.getCustomizedStateConfig(clusterName);
  List<String> listTypesFromZk = configFromZk.getAggregationEnabledTypes();
  Assert.assertEquals(listTypesFromZk.get(0), "mockType1");

  admin.addTypeToCustomizedStateConfig(clusterName, "mockType2");
  admin.addTypeToCustomizedStateConfig(clusterName, "mockType3");
  configFromZk =
      _configAccessor.getCustomizedStateConfig(clusterName);
  listTypesFromZk = configFromZk.getAggregationEnabledTypes();
  Assert.assertEquals(listTypesFromZk.get(0), "mockType1");
  Assert.assertEquals(listTypesFromZk.get(1), "mockType2");
  Assert.assertEquals(listTypesFromZk.get(2), "mockType3");

  admin.removeTypeFromCustomizedStateConfig(clusterName, "mockType1");
  configFromZk =
      _configAccessor.getCustomizedStateConfig(clusterName);
  listTypesFromZk = configFromZk.getAggregationEnabledTypes();
  Assert.assertEquals(listTypesFromZk.get(0), "mockType2");
  Assert.assertEquals(listTypesFromZk.get(1), "mockType3");
}