chore: port a couple of tests from 1x

bigtable-client-core-parent/bigtable-hbase/pom.xml

@@ -194,6 +194,12 @@ limitations under the License.
       <version>${hamcrest.version}</version>
       <scope>test</scope>
     </dependency>
+    <dependency>
+      <groupId>com.google.truth</groupId>
+      <artifactId>truth</artifactId>
+      <version>1.1.2</version>
+      <scope>test</scope>
+    </dependency>
     <dependency>
       <groupId>junit</groupId>
       <artifactId>junit</artifactId>

bigtable-client-core-parent/bigtable-hbase/src/test/java/org/apache/hadoop/hbase/client/TestAbstractBigtableConnection.java

@@ -29,29 +29,38 @@
 import com.google.bigtable.v2.SampleRowKeysResponse;
 import com.google.cloud.bigtable.data.v2.models.KeyOffset;
 import com.google.cloud.bigtable.hbase.AbstractBigtableTable;
+import com.google.cloud.bigtable.hbase.BigtableHBaseVersion;
 import com.google.cloud.bigtable.hbase.BigtableOptionsFactory;
 import com.google.cloud.bigtable.hbase.adapters.SampledRowKeysAdapter;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.Queues;
+import com.google.common.truth.Truth;
+import io.grpc.Metadata;
+import io.grpc.Metadata.Key;
 import io.grpc.Server;
 import io.grpc.ServerBuilder;
+import io.grpc.ServerCall;
+import io.grpc.ServerCall.Listener;
+import io.grpc.ServerCallHandler;
+import io.grpc.ServerInterceptor;
+import io.grpc.internal.GrpcUtil;
 import io.grpc.stub.StreamObserver;
 import java.io.IOException;
 import java.net.ServerSocket;
 import java.util.List;
 import java.util.concurrent.BlockingQueue;
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicReference;
 import org.apache.hadoop.conf.Configuration;
 import org.apache.hadoop.hbase.HRegionInfo;
 import org.apache.hadoop.hbase.HRegionLocation;
 import org.apache.hadoop.hbase.ServerName;
 import org.apache.hadoop.hbase.TableName;
 import org.apache.hadoop.hbase.util.Bytes;
+import org.apache.hadoop.hbase.util.VersionInfo;
 import org.junit.After;
-import org.junit.AfterClass;
 import org.junit.Before;
-import org.junit.BeforeClass;
 import org.junit.Rule;
 import org.junit.Test;
 import org.junit.runner.RunWith;
@@ -77,32 +86,27 @@ public class TestAbstractBigtableConnection {
 
   private static Server server;
   private static int port;
+  private ServerHeaderInterceptor headerInterceptor;
 
   @Mock private AbstractBigtableAdmin mockBigtableAdmin;
 
   @Mock private SampledRowKeysAdapter mockSampledAdapter;
 
   private AbstractBigtableConnection connection;
 
-  @BeforeClass
-  public static void setUpServer() throws IOException {
+  @Before
+  public void setUp() throws IOException {
     try (ServerSocket s = new ServerSocket(0)) {
       port = s.getLocalPort();
     }
-    server = ServerBuilder.forPort(port).addService(fakeDataService).build();
+    headerInterceptor = new ServerHeaderInterceptor();
+    server =
+        ServerBuilder.forPort(port)
+            .intercept(headerInterceptor)
+            .addService(fakeDataService)
+            .build();
     server.start();
-  }
-
-  @AfterClass
-  public static void tearDownServer() throws InterruptedException {
-    if (server != null) {
-      server.shutdownNow();
-      server.awaitTermination();
-    }
-  }
 
-  @Before
-  public void setUp() throws IOException {
     Configuration configuration = new Configuration(false);
     configuration.set(BigtableOptionsFactory.PROJECT_ID_KEY, PROJECT_ID);
     configuration.set(BigtableOptionsFactory.INSTANCE_ID_KEY, INSTANCE_ID);
@@ -113,8 +117,11 @@ public void setUp() throws IOException {
   }
 
   @After
-  public void tearDown() throws IOException {
+  public void tearDown() throws IOException, InterruptedException {
     connection.close();
+    if (server != null) {
+      server.shutdownNow();
+    }
   }
 
   @Test
@@ -168,6 +175,56 @@ public void testConnectionIsClosedWhenAborted() {
     assertTrue(connection.isClosed());
   }
 
+  @Test
+  public void testHeaders() throws IOException {
+    String rowKey = "test-row-key";
+    String value = "mutation-value";
+
+    Table table = connection.getTable(TABLE_NAME);
+    table.put(
+        new Put(Bytes.toBytes(rowKey))
+            .addImmutable(Bytes.toBytes("cf"), Bytes.toBytes("q"), Bytes.toBytes(value)));
+
+    Metadata metadata = headerInterceptor.receivedMetadata.get();
+    String userAgent = metadata.get(Key.of("user-agent", Metadata.ASCII_STRING_MARSHALLER));
+    Truth.assertThat(userAgent).contains("bigtable-" + BigtableHBaseVersion.getVersion());
+    Truth.assertThat(userAgent).contains("hbase-" + VersionInfo.getVersion());
+    Truth.assertThat(userAgent)
+        .contains("grpc-java-netty/" + GrpcUtil.getGrpcBuildVersion().getImplementationVersion());
+    Truth.assertThat(userAgent).contains("jdk-" + System.getProperty("java.specification.version"));
+
+    String xGoogClient =
+        metadata.get(Key.of("x-goog-api-client", Metadata.ASCII_STRING_MARSHALLER));
+    Truth.assertThat(xGoogClient)
+        .contains("gl-java/" + System.getProperty("java.specification.version"));
+    Truth.assertThat(xGoogClient)
+        .contains("grpc/" + GrpcUtil.getGrpcBuildVersion().getImplementationVersion());
+    // Do we need this?
+    // Truth.assertThat(xGoogClient).contains("cbt/" + BigtableHBaseVersion.getVersion());
+
+    // Request must have exactly one of these
+    String requestParams =
+        metadata.get(Key.of("x-goog-request-params", Metadata.ASCII_STRING_MARSHALLER));
+    String resourcePath =
+        metadata.get(Key.of("google-cloud-resource-prefix", Metadata.ASCII_STRING_MARSHALLER));
+
+    if (requestParams != null) {
+      Truth.assertThat(requestParams)
+          .contains(
+              "table_name="
+                  + String.format(
+                      "projects/%s/instances/%s/tables/%s",
+                      PROJECT_ID, INSTANCE_ID, TABLE_NAME.getNameAsString()));
+      Truth.assertThat(resourcePath).isNull();
+    } else {
+      Truth.assertThat(resourcePath)
+          .isEqualTo(
+              String.format(
+                  "projects/%s/instances/%s/tables/%s",
+                  PROJECT_ID, INSTANCE_ID, TABLE_NAME.getNameAsString()));
+    }
+  }
+
   private class TestBigtableConnectionImpl extends AbstractBigtableConnection {
 
     TestBigtableConnectionImpl(Configuration conf) throws IOException {
@@ -221,4 +278,17 @@ public void sampleRowKeys(
       responseObserver.onCompleted();
     }
   }
+
+  private static class ServerHeaderInterceptor implements ServerInterceptor {
+    private AtomicReference<Metadata> receivedMetadata = new AtomicReference<>();
+
+    @Override
+    public <ReqT, RespT> Listener<ReqT> interceptCall(
+        ServerCall<ReqT, RespT> serverCall,
+        Metadata metadata,
+        ServerCallHandler<ReqT, RespT> serverCallHandler) {
+      receivedMetadata.set(metadata);
+      return serverCallHandler.startCall(serverCall, metadata);
+    }
+  }
 }

bigtable-hbase-1.x-parent/bigtable-hbase-1.x/src/test/java/com/google/cloud/bigtable/hbase1_x/BigtableAdminTest.java

@@ -0,0 +1,134 @@
+/*
+ * Copyright 2020 Google LLC
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package com.google.cloud.bigtable.hbase1_x;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertNotNull;
+
+import com.google.bigtable.admin.v2.BigtableTableAdminGrpc;
+import com.google.bigtable.admin.v2.DropRowRangeRequest;
+import com.google.cloud.bigtable.hbase.BigtableConfiguration;
+import com.google.cloud.bigtable.hbase.BigtableOptionsFactory;
+import com.google.protobuf.ByteString;
+import com.google.protobuf.Empty;
+import io.grpc.*;
+import java.io.IOException;
+import java.net.ServerSocket;
+import java.util.concurrent.ArrayBlockingQueue;
+import org.apache.hadoop.conf.Configuration;
+import org.apache.hadoop.hbase.ClusterStatus;
+import org.apache.hadoop.hbase.TableName;
+import org.junit.After;
+import org.junit.Before;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.junit.runners.JUnit4;
+
+@RunWith(JUnit4.class)
+public class BigtableAdminTest {
+  private static final String PROJECT_ID = "fake-project";
+  private static final String INSTANCE_ID = "fake-instance";
+  private static final String TABLE_ID = "fake-table";
+  private static final String TABLE_NAME =
+      String.format("projects/%s/instances/%s/tables/%s", PROJECT_ID, INSTANCE_ID, TABLE_ID);
+
+  private Server fakeBigtableServer;
+  private BigtableConnection connection;
+  private BigtableAdmin admin;
+
+  private ArrayBlockingQueue requestQueue = new ArrayBlockingQueue(1);
+  private ArrayBlockingQueue responseQueue = new ArrayBlockingQueue(1);
+
+  @Before
+  public void setup() throws Exception {
+    final int port;
+
+    try (ServerSocket serverSocket = new ServerSocket(0)) {
+      port = serverSocket.getLocalPort();
+    }
+
+    fakeBigtableServer =
+        ServerBuilder.forPort(port)
+            .intercept(new RequestInterceptor())
+            .addService(new BigtableTableAdminGrpc.BigtableTableAdminImplBase() {})
+            .build();
+    fakeBigtableServer.start();
+
+    Configuration configuration = BigtableConfiguration.configure(PROJECT_ID, INSTANCE_ID);
+    configuration.set(BigtableOptionsFactory.BIGTABLE_EMULATOR_HOST_KEY, "localhost:" + port);
+    connection = new BigtableConnection(configuration);
+    admin = (BigtableAdmin) connection.getAdmin();
+  }
+
+  @After
+  public void teardown() throws IOException {
+    connection.close();
+    fakeBigtableServer.shutdown();
+  }
+
+  @Test
+  public void testDeleteRowRangeByPrefixNonUtf8() throws IOException, InterruptedException {
+    TableName tableName = TableName.valueOf(TABLE_ID);
+    ByteString expectedKey = ByteString.copyFrom(new byte[] {0, 0, 0, (byte) 128});
+
+    DropRowRangeRequest expectedRequest =
+        DropRowRangeRequest.newBuilder().setName(TABLE_NAME).setRowKeyPrefix(expectedKey).build();
+    responseQueue.put(Empty.getDefaultInstance());
+
+    admin.deleteRowRangeByPrefix(tableName, expectedKey.toByteArray());
+
+    assertEquals(expectedRequest, requestQueue.take());
+  }
+
+  private class RequestInterceptor implements ServerInterceptor {
+    @Override
+    public <ReqT, RespT> ServerCall.Listener<ReqT> interceptCall(
+        final ServerCall<ReqT, RespT> call, Metadata headers, ServerCallHandler<ReqT, RespT> next) {
+      return new ServerCall.Listener<ReqT>() {
+        @Override
+        public void onReady() {
+          call.request(1);
+        }
+
+        @Override
+        public void onMessage(ReqT message) {
+          requestQueue.add(message);
+
+          call.sendHeaders(new Metadata());
+          try {
+            @SuppressWarnings("unchecked")
+            RespT response = (RespT) responseQueue.take();
+            call.sendMessage(response);
+            call.close(Status.OK, new Metadata());
+          } catch (InterruptedException e) {
+            call.close(
+                Status.CANCELLED
+                    .withCause(e)
+                    .withDescription("Timed out waiting for mock response"),
+                new Metadata());
+          }
+        }
+      };
+    }
+  }
+
+  @Test
+  public void testGetClusterStatus() throws IOException {
+    // test to verify compatibility between 1x and 2x
+    ClusterStatus status = admin.getClusterStatus();
+    assertNotNull(status);
+  }
+}