From f251b3756f9acc3c82653960b24c6b4c00c2555d Mon Sep 17 00:00:00 2001
From: "J. Liu" <liujiongxin@users.noreply.github.com>
Date: Thu, 26 Oct 2023 10:34:15 -0400
Subject: [PATCH] =?UTF-8?q?test(fix):=20update=20the=20test=20proxy=20to?=
 =?UTF-8?q?=20allow=20enabling=20optional=20features,=E2=80=A6=20(#1976)?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

… also sync the test proto to its latest version

Thank you for opening a Pull Request! Before submitting your PR, there are a few things you can do to make sure it goes smoothly:
- [ ] Make sure to open an issue as a [bug/issue](https://togithub.com/googleapis/java-bigtable/issues/new/choose) before writing your code!  That way we can discuss the change, evaluate designs, and agree on the general idea
- [ ] Ensure the tests and linter pass
- [ ] Code coverage does not decrease (if any source code was changed)
- [ ] Appropriate docs were updated (if necessary)
- [ ] Rollback plan is reviewed and LGTMed

Fixes #<issue_number_goes_here> ☕️

If you write sample code, please follow the [samples format](
https://togithub.com/GoogleCloudPlatform/java-docs-samples/blob/main/SAMPLE_FORMAT.md).
---
 .../cloud/bigtable/testproxy/CbtTestProxy.java    | 15 ++++++++++++++-
 test-proxy/src/main/proto/test_proxy.proto        | 14 ++++++++++++++
 2 files changed, 28 insertions(+), 1 deletion(-)

diff --git a/test-proxy/src/main/java/com/google/cloud/bigtable/testproxy/CbtTestProxy.java b/test-proxy/src/main/java/com/google/cloud/bigtable/testproxy/CbtTestProxy.java
index 5ac1a3c775..3a696dcfe7 100644
--- a/test-proxy/src/main/java/com/google/cloud/bigtable/testproxy/CbtTestProxy.java
+++ b/test-proxy/src/main/java/com/google/cloud/bigtable/testproxy/CbtTestProxy.java
@@ -202,7 +202,7 @@ public synchronized void createClient(
 
     BigtableDataSettings.Builder settingsBuilder =
         BigtableDataSettings.newBuilder()
-            // disable channel refreshing when creating an emulator
+            // Disable channel refreshing when not using the real server
             .setRefreshingChannel(false)
             .setProjectId(request.getProjectId())
             .setInstanceId(request.getInstanceId())
@@ -217,6 +217,19 @@ public synchronized void createClient(
               Durations.toString(request.getPerOperationTimeout())));
     }
 
+    if (request.getOptionalFeatureConfig()
+        == OptionalFeatureConfig.OPTIONAL_FEATURE_CONFIG_ENABLE_ALL) {
+      logger.info("Enabling all the optional features");
+      try {
+        // Exception will be raised if Application Default Credentials is not found.
+        BigtableDataSettings.enableBuiltinMetrics();
+      } catch (IOException e) {
+        responseObserver.onError(Status.INTERNAL.withDescription(e.getMessage()).asException());
+        return;
+      }
+      settingsBuilder.stubSettings().bulkMutateRowsSettings().setServerInitiatedFlowControl(true);
+    }
+
     // Create and store CbtClient for later use
     try {
       if (!request.getDataTarget().equals("emulator")) {
diff --git a/test-proxy/src/main/proto/test_proxy.proto b/test-proxy/src/main/proto/test_proxy.proto
index 551dd4d8b2..e7caef0e7b 100644
--- a/test-proxy/src/main/proto/test_proxy.proto
+++ b/test-proxy/src/main/proto/test_proxy.proto
@@ -26,6 +26,16 @@ option go_package = "./testproxypb";
 option java_multiple_files = true;
 option java_package = "com.google.cloud.bigtable.testproxy";
 
+// A config flag that dictates how the optional features should be enabled
+// during the client creation. The optional features customize how the client
+// interacts with the server, and are defined in
+// https://github.com/googleapis/googleapis/blob/master/google/bigtable/v2/feature_flags.proto
+enum OptionalFeatureConfig {
+  OPTIONAL_FEATURE_CONFIG_DEFAULT = 0;
+
+  OPTIONAL_FEATURE_CONFIG_ENABLE_ALL = 1;
+}
+
 // Request to test proxy service to create a client object.
 message CreateClientRequest {
   // A unique ID associated with the client object to be created.
@@ -52,6 +62,10 @@ message CreateClientRequest {
   // the created client. Otherwise, the default timeout from the client library
   // will be used. Note that the override applies to all the methods.
   google.protobuf.Duration per_operation_timeout = 6;
+
+  // Optional config that dictates how the optional features should be enabled
+  // during the client creation. Please check the enum type's docstring above.
+  OptionalFeatureConfig optional_feature_config = 7;
 }
 
 // Response from test proxy service for CreateClientRequest.