<commit_before>package com.google.cloud.bigtable.hbase;

import com.google.bigtable.admin.cluster.v1.Cluster;
import com.google.bigtable.admin.cluster.v1.CreateClusterRequest;
import com.google.bigtable.admin.cluster.v1.DeleteClusterRequest;
import com.google.bigtable.admin.cluster.v1.GetClusterRequest;
import com.google.bigtable.admin.cluster.v1.ListClustersRequest;
import com.google.bigtable.admin.cluster.v1.ListZonesRequest;
import com.google.bigtable.admin.cluster.v1.ListZonesResponse;
import com.google.bigtable.admin.cluster.v1.Zone;
import com.google.cloud.hadoop.hbase.BigtableClusterAdminClient;
import com.google.cloud.hadoop.hbase.BigtableClusterAdminGrpcClient;
import com.google.common.util.concurrent.UncheckedExecutionException;
import com.google.longrunning.GetOperationRequest;
import com.google.longrunning.Operation;

import org.apache.hadoop.conf.Configuration;
import org.apache.hadoop.hbase.HColumnDescriptor;
import org.apache.hadoop.hbase.HTableDescriptor;
import org.apache.hadoop.hbase.TableName;
import org.apache.hadoop.hbase.client.Admin;
import org.apache.hadoop.hbase.client.BigtableConnection;
import org.apache.hadoop.hbase.client.Connection;
import org.apache.hadoop.hbase.client.Table;
import org.junit.Assert;
import org.junit.Test;

import io.grpc.Status;
import io.grpc.Status.OperationRuntimeException;

import java.io.IOException;
import java.util.Arrays;
import java.util.List;
import java.util.UUID;
import java.util.concurrent.Executors;

/**
 * Tests the Cluster API.
 */
public class TestClusterAPI {

  private static final int MAX_WAIT_SECONDS = 20;
  private static final String TEST_CLUSTER_ID = "test-cluster-api";

  @Test
  public void setup() throws IOException {
    if (!IntegrationTests.isBigtable()) {
      return;
    }
    IntegrationTests.setConfiguration(IntegrationTests.BASE_CONFIGURATION);
    Configuration config = IntegrationTests.getConfiguration();
    BigtableOptions bigtableOptions = BigtableOptionsFactory.fromConfiguration(config);
    BigtableClusterAdminClient client = createClusterAdminStub(bigtableOptions);

    String projectId = bigtableOptions.getProjectId();
    List<Cluster> clusters = getClusters(client, projectId);

    // cleanup any old clusters
    boolean createCluster = true;
    for (Cluster cluster : clusters) {
      if (cluster.getName().contains(TEST_CLUSTER_ID)) {
        dropCluster(client, cluster.getName());
        // createCluster = false;
      }
    }

    List<Zone> zoneList = getZones(client, projectId);
    String fullyQualifiedZoneName = selectZone(zoneList);
    String clusterId = fullyQualifiedZoneName + "/clusters/" + TEST_CLUSTER_ID;

    if (createCluster) {
      Cluster cluster = createACluster(client, fullyQualifiedZoneName, TEST_CLUSTER_ID);
      waitForOperation(client, cluster.getCurrentOperation().getName(), MAX_WAIT_SECONDS);
    }

    Configuration newConfig = newConfiguration(config, clusterId);
    TableName autoDeletedTableName =
        TableName.valueOf("auto-deleted-" + UUID.randomUUID().toString());
    try (Connection connection = new BigtableConnection(newConfig);
        Admin admin = connection.getAdmin()) {
      countTables(admin, 0);
      createTable(admin, autoDeletedTableName);
      countTables(admin, 1);
      TableName tableToDelete = IntegrationTests.newTestTableName();
      createTable(admin, tableToDelete);
      countTables(admin, 2);
      try (Table t = connection.getTable(tableToDelete)) {
        doPutGetDelete(t);
      }
      dropTable(connection, tableToDelete);
      countTables(admin, 1);
    } finally {
      dropCluster(client, clusterId);
    }
  }

  private void countTables(Admin admin, int expectedCount) throws IOException {
    TableName[] tables = admin.listTableNames();
    int actualCount = tables.length;
    Assert.assertEquals(String.format("Got %d tables, expected %d.  The tables: %s", actualCount,
      expectedCount, Arrays.asList(tables)), expectedCount, actualCount);
  }

  private Cluster getCluster(BigtableClusterAdminClient client, String clusterName) {
    GetClusterRequest request = GetClusterRequest.newBuilder().setName(clusterName).build();
    try {
      Cluster response = client.getCluster(request);
      return response;
    } catch (UncheckedExecutionException e) {
      if (e.getCause() != null && e.getCause() instanceof OperationRuntimeException) {
        Status status = ((OperationRuntimeException) e.getCause()).getStatus();
        if (status.getCode() == Status.NOT_FOUND.getCode()) {
          return null;
        }
      }
      e.printStackTrace();
      throw e;
    }
  }

  private void waitForOperation(BigtableClusterAdminClient client, String operationName,
      int maxSeconds) {
    GetOperationRequest request = GetOperationRequest.newBuilder().setName(operationName).build();
    for (int i = 0; i < maxSeconds; i++) {
      try {
        Thread.sleep(1000);
      } catch (InterruptedException e) {
        e.printStackTrace();
      }
      Operation response = client.getOperation(request);
      if (response.getError() != null) {
        return;
      }
    }
    throw new IllegalStateException(String.format(
      "Waited %d seconds and operation was not complete", maxSeconds));
  }

  private BigtableClusterAdminClient createClusterAdminStub(BigtableOptions bigtableOptions)
      throws IOException {
    return BigtableClusterAdminGrpcClient.createClient(
      bigtableOptions.getClusterAdminTransportOptions(), bigtableOptions.getChannelOptions(),
      Executors.newFixedThreadPool(10));
  }

  private List<Zone> getZones(BigtableClusterAdminClient client, String projectId) {
    ListZonesResponse zones =
        client.listZones(ListZonesRequest.newBuilder().setName("/projects/" + projectId).build());
    List<Zone> zoneList = zones.getZonesList();
    Assert.assertTrue("Zones must exist", !zoneList.isEmpty());
    return zoneList;
  }

  private String selectZone(List<Zone> zoneList) {
    int zoneNumber = (int) (zoneList.size() * Math.random());
    return zoneList.get(zoneNumber).getName().replaceFirst("^/", "");
  }

  private Cluster createACluster(BigtableClusterAdminClient client, String zoneName,
      String clusterId) {
    CreateClusterRequest request =
        CreateClusterRequest.newBuilder().setName(zoneName).setClusterId(clusterId)
            .setCluster(Cluster.newBuilder().setServeNodes(1).build()).build();<commit_msg>Removing a '/' prefix from the getZones() call.<commit_after>package com.google.cloud.bigtable.hbase;

import com.google.bigtable.admin.cluster.v1.Cluster;
import com.google.bigtable.admin.cluster.v1.CreateClusterRequest;
import com.google.bigtable.admin.cluster.v1.DeleteClusterRequest;
import com.google.bigtable.admin.cluster.v1.GetClusterRequest;
import com.google.bigtable.admin.cluster.v1.ListClustersRequest;
import com.google.bigtable.admin.cluster.v1.ListZonesRequest;
import com.google.bigtable.admin.cluster.v1.ListZonesResponse;
import com.google.bigtable.admin.cluster.v1.Zone;
import com.google.cloud.hadoop.hbase.BigtableClusterAdminClient;
import com.google.cloud.hadoop.hbase.BigtableClusterAdminGrpcClient;
import com.google.common.util.concurrent.UncheckedExecutionException;
import com.google.longrunning.GetOperationRequest;
import com.google.longrunning.Operation;

import org.apache.hadoop.conf.Configuration;
import org.apache.hadoop.hbase.HColumnDescriptor;
import org.apache.hadoop.hbase.HTableDescriptor;
import org.apache.hadoop.hbase.TableName;
import org.apache.hadoop.hbase.client.Admin;
import org.apache.hadoop.hbase.client.BigtableConnection;
import org.apache.hadoop.hbase.client.Connection;
import org.apache.hadoop.hbase.client.Table;
import org.junit.Assert;
import org.junit.Test;

import io.grpc.Status;
import io.grpc.Status.OperationRuntimeException;

import java.io.IOException;
import java.util.Arrays;
import java.util.List;
import java.util.UUID;
import java.util.concurrent.Executors;

/**
 * Tests the Cluster API.
 */
public class TestClusterAPI {

  private static final int MAX_WAIT_SECONDS = 20;
  private static final String TEST_CLUSTER_ID = "test-cluster-api";

  @Test
  public void setup() throws IOException {
    if (!IntegrationTests.isBigtable()) {
      return;
    }
    IntegrationTests.setConfiguration(IntegrationTests.BASE_CONFIGURATION);
    Configuration config = IntegrationTests.getConfiguration();
    BigtableOptions bigtableOptions = BigtableOptionsFactory.fromConfiguration(config);
    BigtableClusterAdminClient client = createClusterAdminStub(bigtableOptions);

    String projectId = bigtableOptions.getProjectId();
    List<Cluster> clusters = getClusters(client, projectId);

    // cleanup any old clusters
    boolean createCluster = true;
    for (Cluster cluster : clusters) {
      if (cluster.getName().contains(TEST_CLUSTER_ID)) {
        dropCluster(client, cluster.getName());
        // createCluster = false;
      }
    }

    List<Zone> zoneList = getZones(client, projectId);
    String fullyQualifiedZoneName = selectZone(zoneList);
    String clusterId = fullyQualifiedZoneName + "/clusters/" + TEST_CLUSTER_ID;

    if (createCluster) {
      Cluster cluster = createACluster(client, fullyQualifiedZoneName, TEST_CLUSTER_ID);
      waitForOperation(client, cluster.getCurrentOperation().getName(), MAX_WAIT_SECONDS);
    }

    Configuration newConfig = newConfiguration(config, clusterId);
    TableName autoDeletedTableName =
        TableName.valueOf("auto-deleted-" + UUID.randomUUID().toString());
    try (Connection connection = new BigtableConnection(newConfig);
        Admin admin = connection.getAdmin()) {
      countTables(admin, 0);
      createTable(admin, autoDeletedTableName);
      countTables(admin, 1);
      TableName tableToDelete = IntegrationTests.newTestTableName();
      createTable(admin, tableToDelete);
      countTables(admin, 2);
      try (Table t = connection.getTable(tableToDelete)) {
        doPutGetDelete(t);
      }
      dropTable(connection, tableToDelete);
      countTables(admin, 1);
    } finally {
      dropCluster(client, clusterId);
    }
  }

  private void countTables(Admin admin, int expectedCount) throws IOException {
    TableName[] tables = admin.listTableNames();
    int actualCount = tables.length;
    Assert.assertEquals(String.format("Got %d tables, expected %d.  The tables: %s", actualCount,
      expectedCount, Arrays.asList(tables)), expectedCount, actualCount);
  }

  private Cluster getCluster(BigtableClusterAdminClient client, String clusterName) {
    GetClusterRequest request = GetClusterRequest.newBuilder().setName(clusterName).build();
    try {
      Cluster response = client.getCluster(request);
      return response;
    } catch (UncheckedExecutionException e) {
      if (e.getCause() != null && e.getCause() instanceof OperationRuntimeException) {
        Status status = ((OperationRuntimeException) e.getCause()).getStatus();
        if (status.getCode() == Status.NOT_FOUND.getCode()) {
          return null;
        }
      }
      e.printStackTrace();
      throw e;
    }
  }

  private void waitForOperation(BigtableClusterAdminClient client, String operationName,
      int maxSeconds) {
    GetOperationRequest request = GetOperationRequest.newBuilder().setName(operationName).build();
    for (int i = 0; i < maxSeconds; i++) {
      try {
        Thread.sleep(1000);
      } catch (InterruptedException e) {
        e.printStackTrace();
      }
      Operation response = client.getOperation(request);
      if (response.getError() != null) {
        return;
      }
    }
    throw new IllegalStateException(String.format(
      "Waited %d seconds and operation was not complete", maxSeconds));
  }

  private BigtableClusterAdminClient createClusterAdminStub(BigtableOptions bigtableOptions)
      throws IOException {
    return BigtableClusterAdminGrpcClient.createClient(
      bigtableOptions.getClusterAdminTransportOptions(), bigtableOptions.getChannelOptions(),
      Executors.newFixedThreadPool(10));
  }

  private List<Zone> getZones(BigtableClusterAdminClient client, String projectId) {
    ListZonesResponse zones =
        client.listZones(ListZonesRequest.newBuilder().setName("projects/" + projectId).build());
    List<Zone> zoneList = zones.getZonesList();
    Assert.assertTrue("Zones must exist", !zoneList.isEmpty());
    return zoneList;
  }

  private String selectZone(List<Zone> zoneList) {
    int zoneNumber = (int) (zoneList.size() * Math.random());
    return zoneList.get(zoneNumber).getName().replaceFirst("^/", "");
  }

  private Cluster createACluster(BigtableClusterAdminClient client, String zoneName,
      String clusterId) {
    CreateClusterRequest request =
        CreateClusterRequest.newBuilder().setName(zoneName).setClusterId(clusterId)
            .setCluster(Cluster.newBuilder().setServeNodes(1).build()).build();
{'fraction_non_alphanumeric': 0.10664701776157605, 'fraction_numerical': 0.0034127045683704335, 'mean_word_length': 4.312731767614339, 'pattern_counts': {'":': 0, '<': 15, '<?xml version=': 0, '>': 13, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '5662358', 'n_tokens_mistral': 3617, 'n_tokens_neox': 3578, 'n_words': 697}