public static ZKHelixAdmin initHelixClustersForWorkerTest(Properties properties, String route1,
    String route2) throws InterruptedException {
  String zkRoot = properties.getProperty("zkServer");
  Thread.sleep(500);
  ZkClient zkClient = ZkUtils.createZkClient(ZkStarter.DEFAULT_ZK_STR, 1000, 1000);
  zkClient.createPersistent("/ureplicator");
  zkClient.close();
  ZKHelixAdmin helixAdmin = new ZKHelixAdmin(zkRoot);
  String deployment = properties.getProperty("federated.deployment.name");
  String managerHelixClusterName = WorkerUtils.getManagerWorkerHelixClusterName(deployment);
  String controllerHelixClusterName = WorkerUtils.getControllerWorkerHelixClusterName(route1);
  if (StringUtils.isNotBlank(route2)) {
    String controllerHelixClusterName2 = WorkerUtils.getControllerWorkerHelixClusterName(route2);
    HelixSetupUtils.setup(controllerHelixClusterName2, zkRoot, "0");
  }

  HelixSetupUtils.setup(managerHelixClusterName, zkRoot, "0");
  HelixSetupUtils.setup(controllerHelixClusterName, zkRoot, "0");

  return helixAdmin;
}
 

public static void updateRouteWithValidation(String managerHelixClusterName,
    String routeForHelix,
    String instanceId, ZKHelixAdmin helixAdmin, String state, String expectedState)
    throws InterruptedException {
  if (StringUtils.isBlank(expectedState)) {
    expectedState = state;
  }
  IdealState idealState = TestUtils
      .buildManagerWorkerCustomIdealState(routeForHelix, Collections.singletonList(instanceId),
          state);
  helixAdmin.setResourceIdealState(managerHelixClusterName, routeForHelix, idealState);
  Thread.sleep(1000);
  ExternalView externalView = helixAdmin
      .getResourceExternalView(managerHelixClusterName, routeForHelix);
  Assert.assertNotNull(externalView);
  Assert.assertNotNull(externalView.getStateMap("0"));
  LOGGER.info("ExternalView: {}", externalView);
  Assert.assertEquals(externalView.getStateMap("0").get("0"), expectedState);
}
 

public static PistachioClusterInfo getClusterInfo() {
    try {
        String zookeeperConnStr = ConfigurationManager.getConfiguration().getString("Pistachio.ZooKeeper.Server");
        ZKHelixAdmin admin = new ZKHelixAdmin(zookeeperConnStr);
        IdealState idealState = admin.getResourceIdealState("PistachiosCluster", "PistachiosResource");
        PistachioClusterInfo info = new PistachioClusterInfo();
        info.numPartitions = idealState.getNumPartitions();
        info.numReplicas = Integer.parseInt(idealState.getReplicas());
        info.hostList = admin.getInstancesInCluster("PistachiosCluster");

        logger.info("num partitions: {}, num Replicas: {}, hostList: {}.", info.numPartitions,
            info.numReplicas, Joiner.on(",").join(info.hostList.toArray()));

        return info;
    } catch (Exception e) {
        logger.info("error getting cluster info", e);
        return null;
    }
}
 

private static void setConstraints(ZKHelixAdmin admin, int numPartitions) {

        logger.info("pause cluster");
        admin.enableCluster("PistachiosCluster", false);
        // setting partition constraints
        logger.info("setting per partition state transition constraints to 1");
        try {
            for (int constraintId = 0; constraintId < numPartitions; constraintId++) {
                java.util.HashMap<ConstraintAttribute, String>  attributes = new java.util.HashMap<ConstraintAttribute, String>();
                attributes.put(ConstraintAttribute.RESOURCE, "PistachiosResource");
                attributes.put(ConstraintAttribute.PARTITION, "PistachiosResource_"+constraintId);
                logger.info("setting per partition for {} state transition constraints to 1", "PistachiosResource_"+constraintId);
                admin.setConstraint("PistachiosCluster", ConstraintType.STATE_CONSTRAINT, "PistachiosPartitionTransitionConstraint" + constraintId,
                   new ConstraintItem(attributes,"1"));
            }

        } catch(Exception e) {
            logger.info("setting state transition constraints error, roll back and exit", e);
        }
        logger.info("resume cluster");
        admin.enableCluster("PistachiosCluster", true);
  }
 

/**
 * 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();
}
 

public static void setup() {
  admin = new ZKHelixAdmin(ZK_ADDRESS);
  // create cluster
  echo("Creating cluster: " + CLUSTER_NAME);
  admin.addCluster(CLUSTER_NAME, true);

  // Add nodes to the cluster
  echo("Adding " + NUM_NODES + " participants to the cluster");
  for (int i = 0; i < NUM_NODES; i++) {
    admin.addInstance(CLUSTER_NAME, INSTANCE_CONFIG_LIST.get(i));
    echo("\t Added participant: " + INSTANCE_CONFIG_LIST.get(i).getInstanceName());
  }

  // Add a state model
  StateModelDefinition myStateModel = defineStateModel();
  echo("Configuring StateModel: " + "MyStateModel  with 1 Master and 1 Slave");
  admin.addStateModelDef(CLUSTER_NAME, STATE_MODEL_NAME, myStateModel);

  // Add a resource with 6 partitions and 2 replicas
  echo("Adding a resource MyResource: " + "with 6 partitions and 2 replicas");
  admin.addResource(CLUSTER_NAME, RESOURCE_NAME, NUM_PARTITIONS, STATE_MODEL_NAME, "AUTO");
  // this will set up the ideal state, it calculates the preference list for
  // each partition similar to consistent hashing
  admin.rebalance(CLUSTER_NAME, RESOURCE_NAME, NUM_REPLICAS);
}
 

@Deprecated
public ClusterSetup(String zkServerAddress) {
  // If the multi ZK config is enabled, use FederatedZkClient on multi-realm mode
  if (Boolean.parseBoolean(System.getProperty(SystemPropertyKeys.MULTI_ZK_ENABLED))) {
    try {
      _zkClient = new FederatedZkClient(
          new RealmAwareZkClient.RealmAwareZkConnectionConfig.Builder().build(),
          new RealmAwareZkClient.RealmAwareZkClientConfig()
              .setZkSerializer(new ZNRecordSerializer()));
    } catch (IOException | InvalidRoutingDataException | IllegalStateException e) {
      throw new HelixException("Failed to create ConfigAccessor!", e);
    }
  } else {
    _zkClient = SharedZkClientFactory.getInstance()
        .buildZkClient(new HelixZkClient.ZkConnectionConfig(zkServerAddress));
    _zkClient.setZkSerializer(new ZNRecordSerializer());
  }

  _admin = new ZKHelixAdmin(_zkClient);
  _usesExternalZkClient = false;
}
 

@Test()
public void testDisableNode() throws Exception {
  String command =
      "-zkSvr " + ZK_ADDR + " -enableInstance " + CLUSTER_NAME + " " + PARTICIPANT_PREFIX
          + "_12918" + " TestDB TestDB_0 false";
  ClusterSetup.processCommandLineArgs(command.split(" "));
  boolean result =
      ClusterStateVerifier.verifyByPolling(new ClusterStateVerifier.BestPossAndExtViewZkVerifier(
          ZK_ADDR, CLUSTER_NAME));
  Assert.assertTrue(result);

  ZKHelixAdmin tool = new ZKHelixAdmin(_gZkClient);
  tool.enableInstance(CLUSTER_NAME, PARTICIPANT_PREFIX + "_12918", true);

  result =
      ClusterStateVerifier.verifyByPolling(new ClusterStateVerifier.BestPossAndExtViewZkVerifier(
          ZK_ADDR, CLUSTER_NAME));
  Assert.assertTrue(result);

}
 

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

  HelixAdmin firstHelixAdmin = new ZKHelixAdmin.Builder().build();
  HelixAdmin secondHelixAdmin =
      new ZKHelixAdmin.Builder().setRealmAwareZkConnectionConfig(connectionConfig).build();

  try {
    verifyMsdsZkRealm(CLUSTER_ONE, true,
        () -> firstHelixAdmin.enableCluster(CLUSTER_ONE, true));
    verifyMsdsZkRealm(CLUSTER_FOUR, false,
        () -> firstHelixAdmin.enableCluster(CLUSTER_FOUR, true));

    verifyMsdsZkRealm(CLUSTER_FOUR, true,
        () -> secondHelixAdmin.enableCluster(CLUSTER_FOUR, true));
    verifyMsdsZkRealm(CLUSTER_ONE, false,
        () -> secondHelixAdmin.enableCluster(CLUSTER_ONE, true));
  } finally {
    firstHelixAdmin.close();
    secondHelixAdmin.close();
  }
}
 

private static void addInstanceConfig(String instanceName) {
  // add node to cluster if not already added
  ZKHelixAdmin admin = new ZKHelixAdmin(ZK_ADDR);

  InstanceConfig instanceConfig = null;
  try {
    instanceConfig = admin.getInstanceConfig(_clusterName, instanceName);
  } catch (Exception ignored) {
  }
  if (instanceConfig == null) {
    InstanceConfig config = new InstanceConfig(instanceName);
    config.setHostName("localhost");
    config.setInstanceEnabled(true);
    echo("Adding InstanceConfig:" + config);
    admin.addInstance(_clusterName, config);
  }
}
 

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

  ZKHelixAdmin admin = new ZKHelixAdmin(ZK_ADDR);
  admin.addCluster(clusterName, true);
  ConfigAccessor configAccessor = new ConfigAccessor(ZK_ADDR);
  HelixConfigScope scope =
      new HelixConfigScopeBuilder(ConfigScopeProperty.REST).forCluster(clusterName).build();
  Assert.assertNull(configAccessor.getRESTConfig(clusterName));

  RESTConfig restConfig = new RESTConfig(clusterName);
  restConfig.set(RESTConfig.SimpleFields.CUSTOMIZED_HEALTH_URL, "TEST_URL");
  configAccessor.setRESTConfig(clusterName, restConfig);
  Assert.assertEquals(restConfig, configAccessor.getRESTConfig(clusterName));
}
 

public static void waitForExternalViewUpdate(String zkAddress, final String clusterName, long timeoutInMilliseconds) {
  final ZKHelixAdmin helixAdmin = new ZKHelixAdmin(zkAddress);

  List<String> allResourcesInCluster = helixAdmin.getResourcesInCluster(clusterName);
  Set<String> tableAndBrokerResources = new HashSet<>();
  for (String resourceName : allResourcesInCluster) {
    // Only check table resources and broker resource
    if (TableNameBuilder.isTableResource(resourceName) || resourceName
        .equals(CommonConstants.Helix.BROKER_RESOURCE_INSTANCE)) {
      tableAndBrokerResources.add(resourceName);
    }
  }

  StrictMatchExternalViewVerifier verifier =
      new StrictMatchExternalViewVerifier.Builder(clusterName).setZkAddr(zkAddress)
          .setResources(tableAndBrokerResources).build();

  boolean success = verifier.verify(timeoutInMilliseconds);
  if (success) {
    LOGGER.info("Cluster is ready to serve queries");
  }
}
 

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

/**
 * 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!");
}
 

/**
 * Update the resources in the destination cluster with the new IdealState settings.
 * @param destZkString the destination Zookeeper server string.
 * @param destClusterName the destination cluster name.
 * @param dryRun run without actual change.
 */
static void updateResourceIdealState(String destZkString, String destClusterName, boolean dryRun) {
  HelixAdmin destAdmin = new ZKHelixAdmin(destZkString);
  Set<String> destResources = new HashSet<>(destAdmin.getResourcesInCluster(destClusterName));

  for (String resource : destResources) {
    IdealState currentIdealState = destAdmin.getResourceIdealState(destClusterName, resource);
    IdealState newIdealState = buildIdealState(resource, currentIdealState.getPartitionSet());
    if (dryRun) {
      System.out.println("Will update " + resource + " to new ideal state " + newIdealState.toString());
    } else {
      destAdmin.setResourceIdealState(destClusterName, resource, newIdealState);
      System.out.println("Updated the ideal state for resource " + resource);
      destAdmin.rebalance(destClusterName, resource, REPLICA_NUMBER, "", "");
      System.out.println("Rebalanced resource " + resource + " with REPLICA_NUM: " + REPLICA_NUMBER);
    }
  }
}
 

public static void createHelixClusterIfNeeded(String helixClusterName, String zkPath) {
  final HelixAdmin admin = new ZKHelixAdmin(zkPath);

  if (admin.getClusters().contains(helixClusterName)) {
    LOGGER.info("cluster already exist, skipping it.. ********************************************* ");
    return;
  }

  LOGGER.info("Creating a new cluster, as the helix cluster : " + helixClusterName
      + " was not found ********************************************* ");
  admin.addCluster(helixClusterName, false);

  LOGGER.info("Enable mirror maker machines auto join.");
  final HelixConfigScope scope = new HelixConfigScopeBuilder(ConfigScopeProperty.CLUSTER)
      .forCluster(helixClusterName).build();

  final Map<String, String> props = new HashMap<String, String>();
  props.put(ZKHelixManager.ALLOW_PARTICIPANT_AUTO_JOIN, String.valueOf(true));
  props.put(MessageType.STATE_TRANSITION + "." + HelixTaskExecutor.MAX_THREADS,
      String.valueOf(100));

  admin.setConfig(scope, props);

  LOGGER.info("Adding state model definition named : OnlineOffline generated using : "
      + OnlineOfflineStateModel.class.toString()
      + " ********************************************** ");

  // add state model definition
  admin.addStateModelDef(helixClusterName, "OnlineOffline", OnlineOfflineStateModel.build());
  LOGGER.info("New Cluster setup completed... ********************************************** ");
}
 

public static void updateTopicWithValidation(String controllerHelixClusterName, String topicName,
    List<Integer> partitions, List<String> instances, ZKHelixAdmin helixAdmin, String state,
    String expectedState)
    throws InterruptedException {
  Map<String, String> partitionInstanceMap = new HashMap<>();
  if (StringUtils.isBlank(expectedState)) {
    expectedState = state;
  }
  for (int index = 0; index < partitions.size(); index++) {
    partitionInstanceMap
        .put(String.valueOf(partitions.get(index)), instances.get(index % instances.size()));
  }
  IdealState idealState = TestUtils
      .buildControllerWorkerCustomIdealState(topicName, partitionInstanceMap, state);
  LOGGER.info("setResourceIdealState cluster : {}, topic: {} ", controllerHelixClusterName,
      topicName);
  helixAdmin.setResourceIdealState(controllerHelixClusterName, topicName, idealState);
  Thread.sleep(1500);
  ExternalView externalView = helixAdmin
      .getResourceExternalView(controllerHelixClusterName, topicName);
  for (Map.Entry<String, String> entry : partitionInstanceMap.entrySet()) {
    Assert.assertNotNull(externalView);
    Assert.assertNotNull(externalView.getStateMap(entry.getKey()));
    Assert.assertEquals(externalView.getStateMap(entry.getKey()).get(entry.getValue()),
        expectedState);
  }
}
 

private void setUpHelixCluster(String zookeeperQuorum, String clusterName) {
  ZkClient zkClient = ZKClientPool.getZkClient(zookeeperQuorum);
  HelixAdmin helixAdmin = new ZKHelixAdmin(zkClient);
  try {
    if(!ImmutableSet.copyOf(helixAdmin.getClusters()).contains(clusterName)) {
      ClusterSetup helixClusterSetUp = new ClusterSetup(zkClient);
      helixClusterSetUp.addCluster(clusterName, false);
      helixClusterSetUp.setConfig(HelixConfigScope.ConfigScopeProperty.CLUSTER, clusterName,
          "allowParticipantAutoJoin=true");
    }
  } finally {
    zkClient.close();
  }
}
 

private static void cleanup(ZKHelixAdmin admin, ZkClient zkClient, String[] hostList, int numPartitions, int numReplicas, String kafkaTopicPrefix, String kafkaZKPath) {
  try {
      // TODO, delete not supported until 0.8.1, we'll enable it later
      for (int i =0; i<numPartitions; i++) {
          //zkClient = new ZkClient(zkConnect, 30000, 30000, ZKStringSerializer);
          zkClient.deleteRecursive(ZkUtils.getTopicPath(kafkaTopicPrefix + i));
      }
      zkClient.close();
      //ZKHelixAdmin admin = new ZKHelixAdmin(args[1]);
      admin.dropCluster("PistachiosCluster");
    } catch(Exception e) {
        logger.info("error:", e);
    }
      logger.info("cleanup finished succeessfully");
}
 

public long getTotalPartition(String resource) {
    if (totalParition == -1) {
        synchronized(totalParition) {
            if (totalParition == -1) {
                ZKHelixAdmin admin = new ZKHelixAdmin(zkAddress);
                IdealState idealState = admin.getResourceIdealState(helixClusterName, resource);
                totalParition = (long)idealState.getNumPartitions();
            }
        }
    }

    return totalParition;
}
 

@Deprecated
public TaskDriver(RealmAwareZkClient client, ZkBaseDataAccessor<ZNRecord> baseAccessor,
    String clusterName) {
  this(new ZKHelixAdmin(client), new ZKHelixDataAccessor(clusterName, baseAccessor),
      new ZkHelixPropertyStore<>(baseAccessor, PropertyPathBuilder.propertyStore(clusterName),
          null),
      clusterName);
}
 

@BeforeClass
public void beforeClass() throws Exception {
  _numDbs = 1;
  _numNodes = 2;
  _numPartitions = 3;
  _numReplicas = 2;
  _partitionVary = false;
  _admin = new ZKHelixAdmin(_gZkClient);
  _configAccessor = new ConfigAccessor(_gZkClient);
  super.beforeClass();
}
 

private void startAdmin() {
  HelixAdmin admin = new ZKHelixAdmin(ZK_ADDR);

  // create cluster
  System.out.println("Creating cluster: " + _clusterName);
  admin.addCluster(_clusterName, true);

  // add MasterSlave state mode definition
  admin.addStateModelDef(_clusterName, "MasterSlave",
      new StateModelDefinition(generateConfigForMasterSlave()));

  // ideal-state znrecord
  ZNRecord record = new ZNRecord(_resourceName);
  record.setSimpleField("IDEAL_STATE_MODE", "AUTO");
  record.setSimpleField("NUM_PARTITIONS", "1");
  record.setSimpleField("REPLICAS", "2");
  record.setSimpleField("STATE_MODEL_DEF_REF", "MasterSlave");
  record.setListField(_resourceName, Arrays.asList("node1", "node2"));

  admin.setResourceIdealState(_clusterName, _resourceName, new IdealState(record));

  ConstraintItemBuilder builder = new ConstraintItemBuilder();

  // limit one transition message at a time across the entire cluster
  builder.addConstraintAttribute("MESSAGE_TYPE", "STATE_TRANSITION")
      // .addConstraintAttribute("INSTANCE", ".*") // un-comment this line if using instance-level
      // constraint
      .addConstraintAttribute("CONSTRAINT_VALUE", "1");
  admin.setConstraint(_clusterName, ClusterConstraints.ConstraintType.MESSAGE_CONSTRAINT,
      "constraint1", builder.build());
}
 

@Test()
public void testDisablePartition() throws Exception {
  LOG.info("START testDisablePartition() at " + new Date(System.currentTimeMillis()));

  // localhost_12919 is MASTER for TestDB_0
  String command = "--zkSvr " + ZK_ADDR + " --enablePartition false " + CLUSTER_NAME
      + " localhost_12919 TestDB TestDB_0 TestDB_9";
  ClusterSetup.processCommandLineArgs(command.split("\\s+"));
  Map<String, Set<String>> map = new HashMap<>();
  map.put("TestDB_0", TestHelper.setOf("localhost_12919"));
  map.put("TestDB_9", TestHelper.setOf("localhost_12919"));

  boolean result = ClusterStateVerifier.verifyByPolling(
      new ClusterStateVerifier.BestPossAndExtViewZkVerifier(ZK_ADDR, CLUSTER_NAME));
  Assert.assertTrue(result);

  TestHelper.verifyState(CLUSTER_NAME, ZK_ADDR, map, "OFFLINE");

  ZKHelixAdmin tool = new ZKHelixAdmin(_gZkClient);
  tool.enablePartition(true, CLUSTER_NAME, "localhost_12919", "TestDB",
      Collections.singletonList("TestDB_9"));

  result = ClusterStateVerifier.verifyByPolling(
      new ClusterStateVerifier.BestPossAndExtViewZkVerifier(ZK_ADDR, CLUSTER_NAME));
  Assert.assertTrue(result);

  map.clear();
  map.put("TestDB_0", TestHelper.setOf("localhost_12919"));
  TestHelper.verifyState(CLUSTER_NAME, ZK_ADDR, map, "OFFLINE");

  map.clear();
  map.put("TestDB_9", TestHelper.setOf("localhost_12919"));
  TestHelper.verifyState(CLUSTER_NAME, ZK_ADDR, map, "MASTER");

  LOG.info("STOP testDisablePartition() at " + new Date(System.currentTimeMillis()));

}
 

protected void setupInstances(String clusterName, int[] instances) {
  HelixAdmin admin = new ZKHelixAdmin(_gZkClient);
  for (int i = 0; i < instances.length; i++) {
    String instance = "localhost_" + instances[i];
    InstanceConfig instanceConfig = new InstanceConfig(instance);
    instanceConfig.setHostName("localhost");
    instanceConfig.setPort("" + instances[i]);
    instanceConfig.setInstanceEnabled(true);
    admin.addInstance(clusterName, instanceConfig);
  }
}
 

public static void main(String[] args) {
  if (args.length < 1) {
    System.err.println("USAGE: java SetupConsumerCluster zookeeperAddress (e.g. localhost:2181)");
    System.exit(1);
  }

  final String zkAddr = args[0];
  final String clusterName = DEFAULT_CLUSTER_NAME;

  ZkClient zkclient = null;
  try {
    zkclient =
        new ZkClient(zkAddr, ZkClient.DEFAULT_SESSION_TIMEOUT,
            ZkClient.DEFAULT_CONNECTION_TIMEOUT, new ZNRecordSerializer());
    ZKHelixAdmin admin = new ZKHelixAdmin(zkclient);

    // add cluster
    admin.addCluster(clusterName, true);

    // add state model definition
    admin.addStateModelDef(clusterName, DEFAULT_STATE_MODEL,
        new StateModelDefinition(StateModelConfigGenerator.generateConfigForOnlineOffline()));

    // add resource "topic" which has 6 partitions
    String resourceName = DEFAULT_RESOURCE_NAME;
    admin.addResource(clusterName, resourceName, DEFAULT_PARTITION_NUMBER, DEFAULT_STATE_MODEL,
        RebalanceMode.FULL_AUTO.toString());

    admin.rebalance(clusterName, resourceName, 1);

  } finally {
    if (zkclient != null) {
      zkclient.close();
    }
  }
}
 

@Override
public void run() {
  ZkClient zkclient = null;
  try {
    // add node to cluster if not already added
    zkclient =
        new ZkClient(_zkAddr, ZkClient.DEFAULT_SESSION_TIMEOUT,
            ZkClient.DEFAULT_CONNECTION_TIMEOUT, new ZNRecordSerializer());
    ZKHelixAdmin admin = new ZKHelixAdmin(zkclient);

    List<String> nodes = admin.getInstancesInCluster(_clusterName);
    if (!nodes.contains(_instanceName)) {
      InstanceConfig config = new InstanceConfig(_instanceName);
      config.setHostName("localhost");
      config.setInstanceEnabled(true);
      admin.addInstance(_clusterName, config);
    }

    Runtime.getRuntime().addShutdownHook(new Thread() {
      @Override
      public void run() {
        System.out.println("Shutting down " + _instanceName);
        disconnect();
      }
    });

    connect();
  } finally {
    if (zkclient != null) {
      zkclient.close();
    }
  }
}
 

public TaskCluster(String zkAddr, String clusterName) throws Exception {
  _clusterName = clusterName;
  _zkclient =
      new ZkClient(zkAddr, ZkClient.DEFAULT_SESSION_TIMEOUT, ZkClient.DEFAULT_CONNECTION_TIMEOUT,
          new ZNRecordSerializer());
  _admin = new ZKHelixAdmin(_zkclient);
}
 

/**
 * Configure the instance, the configuration of each node is available to
 * other nodes.
 * @param instanceName
 */
private void configureInstance(String instanceName) {
  ZKHelixAdmin helixAdmin = new ZKHelixAdmin(zkAddress);

  List<String> instancesInCluster = helixAdmin.getInstancesInCluster(clusterName);
  if (instancesInCluster == null || !instancesInCluster.contains(instanceName)) {
    InstanceConfig config = new InstanceConfig(instanceName);
    config.setHostName("localhost");
    config.setPort("12000");
    helixAdmin.addInstance(clusterName, config);
  }
}
 

private void init() {
  LOGGER.info("Trying to connect to " + _zkAddress + " cluster " + _clusterName);
  _helixAdmin = new ZKHelixAdmin(_zkAddress);
  ZNRecordSerializer serializer = new ZNRecordSerializer();
  String path = PropertyPathConfig.getPath(PropertyType.PROPERTYSTORE, _clusterName);
  _propertyStore = new ZkHelixPropertyStore<>(_zkAddress, serializer, path);
}