org.apache.helix.HelixAdmin Java Examples

@Test
public void testDisableResource() {
  String className = TestHelper.getTestClassName();
  String methodName = TestHelper.getTestMethodName();
  String clusterName = className + "_" + methodName;
  System.out.println("START " + clusterName + " at " + new Date(System.currentTimeMillis()));
  HelixAdmin admin = new ZKHelixAdmin(_gZkClient);
  admin.addCluster(clusterName, true);
  Assert.assertTrue(ZKUtil.isClusterSetup(clusterName, _gZkClient), "Cluster should be setup");
  String resourceName = "TestDB";
  admin.addStateModelDef(clusterName, "MasterSlave",
      new StateModelDefinition(StateModelConfigGenerator.generateConfigForMasterSlave()));
  admin.addResource(clusterName, resourceName, 4, "MasterSlave");
  admin.enableResource(clusterName, resourceName, false);
  BaseDataAccessor<ZNRecord> baseAccessor = new ZkBaseDataAccessor<>(_gZkClient);
  HelixDataAccessor accessor = new ZKHelixDataAccessor(clusterName, baseAccessor);
  PropertyKey.Builder keyBuilder = accessor.keyBuilder();
  IdealState idealState = accessor.getProperty(keyBuilder.idealStates(resourceName));
  Assert.assertFalse(idealState.isEnabled());
  admin.enableResource(clusterName, resourceName, true);
  idealState = accessor.getProperty(keyBuilder.idealStates(resourceName));
  Assert.assertTrue(idealState.isEnabled());

  admin.dropCluster(clusterName);
  System.out.println("END " + clusterName + " at " + new Date(System.currentTimeMillis()));
}
 

/**
 * Create a helix cluster with given information.
 * @param destZkString the cluster's zk string
 * @param destClusterName the cluster's name
 */
static void createCluster(String destZkString, String destClusterName) {
  HelixZkClient destZkClient = getHelixZkClient(destZkString);
  HelixAdmin destAdmin = new ZKHelixAdmin(destZkClient);
  if (ZKUtil.isClusterSetup(destClusterName, destZkClient)) {
    errorAndExit("Failed to create cluster because " + destClusterName + " already exist.");
  }
  ClusterSetup clusterSetup = new ClusterSetup.Builder().setZkAddress(destZkString).build();
  clusterSetup.addCluster(destClusterName, true);

  // set ALLOW_PARTICIPANT_AUTO_JOIN
  HelixConfigScope configScope = new HelixConfigScopeBuilder(HelixConfigScope.ConfigScopeProperty.CLUSTER).
      forCluster(destClusterName).build();
  Map<String, String> helixClusterProperties = new HashMap<>();
  helixClusterProperties.put(ZKHelixManager.ALLOW_PARTICIPANT_AUTO_JOIN, String.valueOf(true));
  destAdmin.setConfig(configScope, helixClusterProperties);

  setClusterConfig(destZkClient, destClusterName, false);
  System.out.println("Cluster " + destClusterName + " is created successfully!");
}
 

/**
 * Add new state model def to ambry cluster in enabled datacenter(s).
 */
private void addStateModelDef() {
  for (Map.Entry<String, HelixAdmin> entry : adminForDc.entrySet()) {
    // Add a cluster entry in every enabled DC
    String dcName = entry.getKey();
    HelixAdmin admin = entry.getValue();
    boolean isClusterPresent = zkClientForDc.get(dcName) == null ? admin.getClusters().contains(clusterName)
        : ZKUtil.isClusterSetup(clusterName, zkClientForDc.get(dcName));
    if (!isClusterPresent) {
      throw new IllegalStateException("Cluster " + clusterName + " in " + dcName + " doesn't exist!");
    }
    if (!admin.getStateModelDefs(clusterName).contains(stateModelDef)) {
      info("Adding state model def {} in {} for cluster {}", stateModelDef, dcName, clusterName);
      admin.addStateModelDef(clusterName, stateModelDef, getStateModelDefinition(stateModelDef));
    } else {
      info("{} in {} already has state model def {}, skip adding operation", clusterName, dcName, stateModelDef);
    }
  }
}
 

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

  HelixAdmin admin = getHelixAdmin();
  try {
    admin.removeCustomizedStateConfig(clusterId);
  } catch (Exception ex) {
    LOG.error(
        "Cannot remove CustomizedStateConfig from cluster: {}, Exception: {}",
        clusterId, ex);
    return serverError(ex);
  }

  return OK();
}
 

/**
 * Reset the partition on specific node.
 */
private void resetPartition() {
  if (adminForDc.size() != 1) {
    throw new IllegalStateException("The dc count is not 1 for resetting partition operation");
  }
  HelixAdmin helixAdmin = adminForDc.values().iterator().next();
  String instanceName;
  if (portNum == null) {
    Optional<DataNodeId> optionalDataNode =
        staticClusterMap.getDataNodeIds().stream().filter(node -> node.getHostname().equals(hostName)).findFirst();
    if (!optionalDataNode.isPresent()) {
      throw new IllegalStateException("Host " + hostName + " is not found in static clustermap");
    }
    DataNodeId dataNodeId = optionalDataNode.get();
    instanceName = getInstanceName(dataNodeId);
  } else {
    instanceName = ClusterMapUtils.getInstanceName(hostName, portNum);
  }
  String resourceName = getResourceNameOfPartition(helixAdmin, clusterName, partitionName);
  info("Resetting partition {} under resource {} on node {}", partitionName, resourceName, hostName);
  helixAdmin.resetPartition(clusterName, instanceName, resourceName, Collections.singletonList(partitionName));
  partitionsReset.getAndIncrement();
}
 

@GET
@Path("/cluster/configs")
@Produces(MediaType.APPLICATION_JSON)
@ApiOperation(value = "List cluster configurations", notes = "List cluster level configurations")
@ApiResponses(value = {@ApiResponse(code = 200, message = "Success"), @ApiResponse(code = 500, message = "Internal server error")})
public String listClusterConfigs() {
  HelixAdmin helixAdmin = pinotHelixResourceManager.getHelixAdmin();
  HelixConfigScope configScope = new HelixConfigScopeBuilder(HelixConfigScope.ConfigScopeProperty.CLUSTER)
      .forCluster(pinotHelixResourceManager.getHelixClusterName()).build();
  List<String> configKeys = helixAdmin.getConfigKeys(configScope);
  ObjectNode ret = JsonUtils.newObjectNode();
  Map<String, String> configs = helixAdmin.getConfig(configScope, configKeys);
  for (String key : configs.keySet()) {
    ret.put(key, configs.get(key));
  }
  return ret.toString();
}
 

private static void setupHelixClusterIfNeeded(String helixClusterName, String zkPath) {
  HelixAdmin admin = new ZKHelixAdmin(zkPath);
  if (admin.getClusters().contains(helixClusterName)) {
    LOGGER.info("Helix cluster: {} already exists", helixClusterName);
  } else {
    LOGGER.info("Creating a new Helix cluster: {}", helixClusterName);
    admin.addCluster(helixClusterName, false);
    // Enable Auto-Join for the cluster
    HelixConfigScope configScope =
        new HelixConfigScopeBuilder(ConfigScopeProperty.CLUSTER).forCluster(helixClusterName).build();
    Map<String, String> configMap = new HashMap<>();
    configMap.put(ZKHelixManager.ALLOW_PARTICIPANT_AUTO_JOIN, Boolean.toString(true));
    configMap.put(ENABLE_CASE_INSENSITIVE_KEY, Boolean.toString(false));
    configMap.put(DEFAULT_HYPERLOGLOG_LOG2M_KEY, Integer.toString(DEFAULT_HYPERLOGLOG_LOG2M));
    configMap.put(CommonConstants.Broker.CONFIG_OF_ENABLE_QUERY_LIMIT_OVERRIDE, Boolean.toString(false));
    admin.setConfig(configScope, configMap);
    LOGGER.info("New Helix cluster: {} created", helixClusterName);
  }
}
 

/**
 * A method to handle failures joining Helix cluster. The method will perform the following steps before attempting
 * to re-join the cluster:
 * <li>
 *   <ul>Disconnect from Helix cluster, which would close any open clients</ul>
 *   <ul>Drop instance from Helix cluster, to remove any partial instance structure from Helix</ul>
 *   <ul>Re-construct helix manager instances, used to re-join the cluster</ul>
 * </li>
 */
private void onClusterJoinFailure() {
  logger.warn("Disconnecting Helix manager..");
  disconnectHelixManager();

  HelixAdmin admin = new ZKHelixAdmin(clusterConfig.getString(GobblinClusterConfigurationKeys.ZK_CONNECTION_STRING_KEY));
  //Drop the helix Instance
  logger.warn("Dropping instance: {} from cluster: {}", helixInstanceName, clusterName);
  HelixUtils.dropInstanceIfExists(admin, clusterName, helixInstanceName);

  if (this.taskDriverHelixManager.isPresent()) {
    String taskDriverCluster = clusterConfig.getString(GobblinClusterConfigurationKeys.TASK_DRIVER_CLUSTER_NAME_KEY);
    logger.warn("Dropping instance: {} from task driver cluster: {}", helixInstanceName, taskDriverCluster);
    HelixUtils.dropInstanceIfExists(admin, clusterName, helixInstanceName);
  }
  admin.close();

  logger.warn("Reinitializing Helix manager..");
  initHelixManager();
}
 

private static HelixManager getMockHelixManager(Config config) {
  HelixManager helixManager = Mockito.mock(HelixManager.class);
  HelixAdmin helixAdmin = Mockito.mock(HelixAdmin.class);
  HelixDataAccessor helixDataAccessor = Mockito.mock(HelixDataAccessor.class);
  PropertyKey propertyKey = Mockito.mock(PropertyKey.class);
  PropertyKey.Builder propertyKeyBuilder = Mockito.mock(PropertyKey.Builder.class);

  Mockito.when(helixManager.getInstanceName()).thenReturn("helixInstance1");
  Mockito.when(helixManager.getClusterName()).thenReturn(config.getString(GobblinClusterConfigurationKeys.HELIX_CLUSTER_NAME_KEY));
  Mockito.doNothing().when(helixAdmin).enableInstance(Mockito.anyString(), Mockito.anyString(), Mockito.anyBoolean());
  Mockito.when(helixManager.getHelixDataAccessor()).thenReturn(helixDataAccessor);
  Mockito.when(helixDataAccessor.keyBuilder()).thenReturn(propertyKeyBuilder);
  Mockito.when(propertyKeyBuilder.liveInstance(Mockito.anyString())).thenReturn(propertyKey);
  Mockito.when(helixDataAccessor.getProperty(propertyKey)).thenReturn(null);

  return helixManager;
}
 

private Set<String> getCurrentPartitionsOnInstance(String cluster, String dbName,
    String instance) {
  HelixAdmin admin = _gSetupTool.getClusterManagementTool();
  Set<String> partitionSet = new HashSet<>();

  ExternalView ev = admin.getResourceExternalView(cluster, dbName);
  for (String partition : ev.getRecord().getMapFields().keySet()) {
    Map<String, String> assignments = ev.getRecord().getMapField(partition);
    for (String ins : assignments.keySet()) {
      if (ins.equals(instance)) {
        partitionSet.add(partition);
      }
    }
  }
  return partitionSet;
}
 

/**
 * Test enabledWagedRebalance by checking the rebalancer class name changed.
 */
@Test
public void testEnableWagedRebalance() {
  String className = TestHelper.getTestClassName();
  String methodName = TestHelper.getTestMethodName();
  String clusterName = className + "_" + methodName;
  String testResourcePrefix = "TestResource";
  HelixAdmin admin = new ZKHelixAdmin(_gZkClient);
  admin.addCluster(clusterName, true);
  admin.addStateModelDef(clusterName, "MasterSlave", new MasterSlaveSMD());

  // Add an IdealState
  IdealState idealState = new IdealState(testResourcePrefix);
  idealState.setNumPartitions(3);
  idealState.setStateModelDefRef("MasterSlave");
  idealState.setRebalanceMode(IdealState.RebalanceMode.FULL_AUTO);
  admin.addResource(clusterName, testResourcePrefix, idealState);

  admin.enableWagedRebalance(clusterName, Collections.singletonList(testResourcePrefix));
  IdealState is = admin.getResourceIdealState(clusterName, testResourcePrefix);
  Assert.assertEquals(is.getRebalancerClassName(), WagedRebalancer.class.getName());
}
 

protected void addFakeBrokerInstanceToAutoJoinHelixCluster(String instanceId, boolean isSingleTenant)
    throws Exception {
  HelixManager helixManager = HelixManagerFactory
      .getZKHelixManager(getHelixClusterName(), instanceId, InstanceType.PARTICIPANT, ZkStarter.DEFAULT_ZK_STR);
  helixManager.getStateMachineEngine()
      .registerStateModelFactory(FakeBrokerResourceOnlineOfflineStateModelFactory.STATE_MODEL_DEF,
          FakeBrokerResourceOnlineOfflineStateModelFactory.FACTORY_INSTANCE);
  helixManager.connect();
  HelixAdmin helixAdmin = helixManager.getClusterManagmentTool();
  if (isSingleTenant) {
    helixAdmin.addInstanceTag(getHelixClusterName(), instanceId, TagNameUtils.getBrokerTagForTenant(null));
  } else {
    helixAdmin.addInstanceTag(getHelixClusterName(), instanceId, UNTAGGED_BROKER_INSTANCE);
  }
  _fakeInstanceHelixManagers.add(helixManager);
}
 

protected void addFakeServerInstanceToAutoJoinHelixCluster(String instanceId, boolean isSingleTenant, int adminPort)
    throws Exception {
  HelixManager helixManager = HelixManagerFactory
      .getZKHelixManager(getHelixClusterName(), instanceId, InstanceType.PARTICIPANT, ZkStarter.DEFAULT_ZK_STR);
  helixManager.getStateMachineEngine()
      .registerStateModelFactory(FakeSegmentOnlineOfflineStateModelFactory.STATE_MODEL_DEF,
          FakeSegmentOnlineOfflineStateModelFactory.FACTORY_INSTANCE);
  helixManager.connect();
  HelixAdmin helixAdmin = helixManager.getClusterManagmentTool();
  if (isSingleTenant) {
    helixAdmin.addInstanceTag(getHelixClusterName(), instanceId, TagNameUtils.getOfflineTagForTenant(null));
    helixAdmin.addInstanceTag(getHelixClusterName(), instanceId, TagNameUtils.getRealtimeTagForTenant(null));
  } else {
    helixAdmin.addInstanceTag(getHelixClusterName(), instanceId, UNTAGGED_SERVER_INSTANCE);
  }
  HelixConfigScope configScope =
      new HelixConfigScopeBuilder(HelixConfigScope.ConfigScopeProperty.PARTICIPANT, getHelixClusterName())
          .forParticipant(instanceId).build();
  helixAdmin.setConfig(configScope, Collections.singletonMap(ADMIN_PORT_KEY, Integer.toString(adminPort)));
  _fakeInstanceHelixManagers.add(helixManager);
}
 

public static Set<TopicPartition> getUnassignedPartitions(HelixManager helixManager) {
  Set<TopicPartition> unassignedPartitions = new HashSet<TopicPartition>();
  HelixAdmin helixAdmin = helixManager.getClusterManagmentTool();
  String helixClusterName = helixManager.getClusterName();
  for (String topic : helixAdmin.getResourcesInCluster(helixClusterName)) {
    IdealState is = helixAdmin.getResourceIdealState(helixClusterName, topic);
    int numPartitions = is.getNumPartitions();
    for (int partition = 0; partition < numPartitions; ++partition) {
      if (is.getInstanceSet(Integer.toString(partition)).isEmpty()) {
        TopicPartition tpi = new TopicPartition(topic, partition);
        unassignedPartitions.add(tpi);
      }
    }
  }
  return unassignedPartitions;
}
 

private void checkInstanceState(HelixAdmin helixAdmin) {
  String expectedInstanceState = "MASTER";
  TestUtils.waitForCondition(aVoid -> {
    ExternalView leadControllerResourceExternalView =
        helixAdmin.getResourceExternalView(getHelixClusterName(), Helix.LEAD_CONTROLLER_RESOURCE_NAME);
    Set<String> partitionSet = leadControllerResourceExternalView.getPartitionSet();
    if (partitionSet.size() != Helix.NUMBER_OF_PARTITIONS_IN_LEAD_CONTROLLER_RESOURCE) {
      return false;
    }
    for (String partition : partitionSet) {
      Map<String, String> stateMap = leadControllerResourceExternalView.getStateMap(partition);
      if (stateMap.size() != 1 || !stateMap.values().contains(expectedInstanceState)) {
        return false;
      }
    }
    return true;
  }, TIMEOUT_IN_MS, "Failed to pick only one instance as: " + expectedInstanceState);
}
 

@DELETE
public Response deleteInstance(@PathParam("clusterId") String clusterId,
    @PathParam("instanceName") String instanceName) {
  HelixAdmin admin = getHelixAdmin();
  try {
    InstanceConfig instanceConfig = admin.getInstanceConfig(clusterId, instanceName);
    admin.dropInstance(clusterId, instanceConfig);
  } catch (HelixException e) {
    return badRequest(e.getMessage());
  }

  return OK();
}
 

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

  HelixAdmin admin = getHelixAdmin();
  ZNRecord record;
  try {
    record = toZNRecord(content);
  } catch (IOException e) {
    return badRequest("Input is not a vaild ZNRecord!");
  }

  try {
    CustomizedStateConfig customizedStateConfig =
        new CustomizedStateConfig.Builder(record).build();
    admin.addCustomizedStateConfig(clusterId, customizedStateConfig);
  } catch (Exception ex) {
    LOG.error("Cannot add CustomizedStateConfig to cluster: {} Exception: {}",
        clusterId, ex);
    return serverError(ex);
  }

  return OK();
}
 

@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);
}
 

private void testAllFailed(List<String> segments) {
  HelixAdmin helixAdmin = makeHelixAdmin();
  ZkHelixPropertyStore<ZNRecord> propertyStore = makePropertyStore();
  FakeDeletionManager deletionManager = new FakeDeletionManager(helixAdmin, propertyStore);
  deletionManager.deleteSegmentsFromPropertyStoreAndLocal(tableName, segments);

  Assert.assertTrue(deletionManager.segmentsToRetry.containsAll(segments));
  Assert.assertEquals(deletionManager.segmentsToRetry.size(), segments.size());
  Assert.assertEquals(deletionManager.segmentsRemovedFromStore.size(), 0);
}
 

public static void setupPinotCluster(String helixClusterName, String zkPath, boolean isUpdateStateModel,
    boolean enableBatchMessageMode, String leadControllerResourceRebalanceStrategy) {
  HelixZkClient zkClient = null;
  try {
    zkClient = SharedZkClientFactory.getInstance().buildZkClient(new HelixZkClient.ZkConnectionConfig(zkPath),
        new HelixZkClient.ZkClientConfig().setZkSerializer(new ZNRecordSerializer())
            .setConnectInitTimeout(TimeUnit.SECONDS.toMillis(ZkClient.DEFAULT_CONNECT_TIMEOUT_SEC)));
    zkClient.waitUntilConnected(ZkClient.DEFAULT_CONNECT_TIMEOUT_SEC, TimeUnit.SECONDS);
    HelixAdmin helixAdmin = new ZKHelixAdmin(zkClient);
    HelixDataAccessor helixDataAccessor =
        new ZKHelixDataAccessor(helixClusterName, new ZkBaseDataAccessor<>(zkClient));
    ConfigAccessor configAccessor = new ConfigAccessor(zkClient);

    Preconditions.checkState(helixAdmin.getClusters().contains(helixClusterName),
        String.format("Helix cluster: %s hasn't been set up", helixClusterName));

    // Add segment state model definition if needed
    addSegmentStateModelDefinitionIfNeeded(helixClusterName, helixAdmin, helixDataAccessor, isUpdateStateModel);

    // Add broker resource if needed
    createBrokerResourceIfNeeded(helixClusterName, helixAdmin, enableBatchMessageMode);

    // Add lead controller resource if needed
    createLeadControllerResourceIfNeeded(helixClusterName, helixAdmin, configAccessor, enableBatchMessageMode,
        leadControllerResourceRebalanceStrategy);
  } finally {
    if (zkClient != null) {
      zkClient.close();
    }
  }
}
 

@DELETE
@Path("{clusterId}/cloudconfig")
public Response deleteCloudConfig(@PathParam("clusterId") String clusterId) {
  HelixAdmin admin = getHelixAdmin();
  admin.removeCloudConfig(clusterId);
  return OK();
}
 

private static void addSegmentStateModelDefinitionIfNeeded(String helixClusterName, HelixAdmin helixAdmin,
    HelixDataAccessor helixDataAccessor, boolean isUpdateStateModel) {
  String segmentStateModelName =
      PinotHelixSegmentOnlineOfflineStateModelGenerator.PINOT_SEGMENT_ONLINE_OFFLINE_STATE_MODEL;
  StateModelDefinition stateModelDefinition = helixAdmin.getStateModelDef(helixClusterName, segmentStateModelName);
  if (stateModelDefinition == null || isUpdateStateModel) {
    if (stateModelDefinition == null) {
      LOGGER.info("Adding state model: {} with CONSUMING state", segmentStateModelName);
    } else {
      LOGGER.info("Updating state model: {} to contain CONSUMING state", segmentStateModelName);
    }
    helixDataAccessor
        .createStateModelDef(PinotHelixSegmentOnlineOfflineStateModelGenerator.generatePinotStateModelDefinition());
  }
}
 

@Test
public void testRebuildBrokerResourceWhenBrokerAdded()
    throws Exception {
  // Check that the first table we added doesn't need to be rebuilt(case where ideal state brokers and brokers in broker resource are the same.
  String partitionName = _offlineTableConfig.getTableName();
  HelixAdmin helixAdmin = _helixManager.getClusterManagmentTool();

  IdealState idealState = HelixHelper.getBrokerIdealStates(helixAdmin, getHelixClusterName());
  // Ensure that the broker resource is not rebuilt.
  Assert.assertTrue(idealState.getInstanceSet(partitionName)
      .equals(_helixResourceManager.getAllInstancesForBrokerTenant(TagNameUtils.DEFAULT_TENANT_NAME)));
  _helixResourceManager.rebuildBrokerResourceFromHelixTags(partitionName);

  // Add another table that needs to be rebuilt
  TableConfig offlineTableConfigTwo = new TableConfigBuilder(TableType.OFFLINE).setTableName(TEST_TABLE_TWO).build();
  _helixResourceManager.addTable(offlineTableConfigTwo);
  String partitionNameTwo = offlineTableConfigTwo.getTableName();

  // Add a new broker manually such that the ideal state is not updated and ensure that rebuild broker resource is called
  final String brokerId = "Broker_localhost_2";
  InstanceConfig instanceConfig = new InstanceConfig(brokerId);
  instanceConfig.setInstanceEnabled(true);
  instanceConfig.setHostName("Broker_localhost");
  instanceConfig.setPort("2");
  helixAdmin.addInstance(getHelixClusterName(), instanceConfig);
  helixAdmin.addInstanceTag(getHelixClusterName(), instanceConfig.getInstanceName(),
      TagNameUtils.getBrokerTagForTenant(TagNameUtils.DEFAULT_TENANT_NAME));
  idealState = HelixHelper.getBrokerIdealStates(helixAdmin, getHelixClusterName());
  // Assert that the two don't equal before the call to rebuild the broker resource.
  Assert.assertTrue(!idealState.getInstanceSet(partitionNameTwo)
      .equals(_helixResourceManager.getAllInstancesForBrokerTenant(TagNameUtils.DEFAULT_TENANT_NAME)));
  _helixResourceManager.rebuildBrokerResourceFromHelixTags(partitionNameTwo);
  idealState = HelixHelper.getBrokerIdealStates(helixAdmin, getHelixClusterName());
  // Assert that the two do equal after being rebuilt.
  Assert.assertTrue(idealState.getInstanceSet(partitionNameTwo)
      .equals(_helixResourceManager.getAllInstancesForBrokerTenant(TagNameUtils.DEFAULT_TENANT_NAME)));
}
 

public static void deleteResourcePropertyFromHelix(HelixAdmin admin, String clusterName, String resourceName,
    String configKey) {
  final List<String> keys = new ArrayList<String>();
  keys.add(configKey);

  final HelixConfigScope scope = getResourceScopeFor(clusterName, resourceName);
  admin.removeConfig(scope, keys);
}
 

/**
 * A method to verify resources and partitions in src cluster and dest cluster are same.
 */
private boolean isSrcDestSync(String srcZkString, String srcClusterName, String destZkString,
    String destClusterName) {

  HelixAdmin srcAdmin = new ZKHelixAdmin(srcZkString);
  Set<String> srcResources = new HashSet<>(srcAdmin.getResourcesInCluster(srcClusterName));
  HelixAdmin destAdmin = new ZKHelixAdmin(destZkString);
  Set<String> destResources = new HashSet<>(destAdmin.getResourcesInCluster(destClusterName));

  for (String resource : srcResources) {
    if (HelixVcrPopulateTool.ignoreResourceKeyWords.stream().anyMatch(resource::contains)) {
      System.out.println("Resource " + resource + " from src cluster is ignored");
      continue;
    }
    if (destResources.contains(resource)) {
      // check if every partition exist.
      Set<String> srcPartitions = srcAdmin.getResourceIdealState(srcClusterName, resource).getPartitionSet();
      Set<String> destPartitions = destAdmin.getResourceIdealState(destClusterName, resource).getPartitionSet();
      for (String partition : srcPartitions) {
        if (!destPartitions.contains(partition)) {
          return false;
        }
      }
    } else {
      return false;
    }
  }
  return true;
}
 

private List<ExternalView> getExternalViews() {
  String clusterName = helixManager.getClusterName();
  List<ExternalView> externalViewList = new ArrayList<ExternalView>();
  HelixAdmin helixAdmin = helixManager.getClusterManagmentTool();
  List<String> resourcesInCluster = helixAdmin.getResourcesInCluster(clusterName);
  for (String resourceName : resourcesInCluster) {
    ExternalView ev =
        helixManager.getClusterManagmentTool().getResourceExternalView(clusterName, resourceName);
    if (ev != null) {
      externalViewList.add(ev);
    }
  }

  return externalViewList;
}
 

@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");
}
 

@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);
}
 

private void assignIdealStates(HelixManager helixManager,
    Map<String, IdealState> idealStatesFromAssignment) {
  HelixAdmin helixAdmin = helixManager.getClusterManagmentTool();
  String helixClusterName = helixManager.getClusterName();
  for (String topic : idealStatesFromAssignment.keySet()) {
    IdealState idealState = idealStatesFromAssignment.get(topic);
    helixAdmin.setResourceIdealState(helixClusterName, topic, idealState);
  }
}
 

@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");
}