From f77465e094ca9b00fc2eb6882a69b9eb9dfd8edb Mon Sep 17 00:00:00 2001
From: Gaole Meng <gaole@google.com>
Date: Tue, 28 Nov 2023 10:49:47 -0800
Subject: [PATCH] fix: an atempt to solve test failure in nightly build (#2330)

---
 .../bigquery/storage/v1/StreamWriter.java     | 30 ++++++++++---------
 .../it/ITBigQueryWriteManualClientTest.java   |  9 ++++++
 2 files changed, 25 insertions(+), 14 deletions(-)

diff --git a/google-cloud-bigquerystorage/src/main/java/com/google/cloud/bigquery/storage/v1/StreamWriter.java b/google-cloud-bigquerystorage/src/main/java/com/google/cloud/bigquery/storage/v1/StreamWriter.java
index b9d16ca84a..e429678d8d 100644
--- a/google-cloud-bigquerystorage/src/main/java/com/google/cloud/bigquery/storage/v1/StreamWriter.java
+++ b/google-cloud-bigquerystorage/src/main/java/com/google/cloud/bigquery/storage/v1/StreamWriter.java
@@ -371,22 +371,24 @@ static BigQueryWriteSettings getBigQueryWriteSettings(Builder builder) throws IO
 
   // Validate whether the fetched connection pool matched certain properties.
   private void validateFetchedConnectonPool(StreamWriter.Builder builder) {
-    String paramsValidatedFailed = "";
-    if (!Objects.equals(
-        this.singleConnectionOrConnectionPool.connectionWorkerPool().getTraceId(),
-        builder.traceId)) {
-      paramsValidatedFailed = "Trace id";
-    } else if (!Objects.equals(
-        this.singleConnectionOrConnectionPool.connectionWorkerPool().limitExceededBehavior(),
-        builder.limitExceededBehavior)) {
-      paramsValidatedFailed = "Limit Exceeds Behavior";
-    }
-
-    if (!paramsValidatedFailed.isEmpty()) {
+    String storedTraceId =
+        this.singleConnectionOrConnectionPool.connectionWorkerPool().getTraceId();
+    if (!Objects.equals(storedTraceId, builder.traceId)) {
       throw new IllegalArgumentException(
           String.format(
-              "%s used for the same connection pool for the same location must be the same!",
-              paramsValidatedFailed));
+              "Trace id used for the same connection pool for the same location must be the same, "
+                  + "however stored trace id is %s, and expected trace id is %s.",
+              storedTraceId, builder.traceId));
+    }
+    FlowController.LimitExceededBehavior storedLimitExceededBehavior =
+        singleConnectionOrConnectionPool.connectionWorkerPool().limitExceededBehavior();
+    if (!Objects.equals(storedLimitExceededBehavior, builder.limitExceededBehavior)) {
+      throw new IllegalArgumentException(
+          String.format(
+              "Limit exceeded behavior setting used for the same connection pool for the same "
+                  + "location must be the same, however stored value is %s, and expected "
+                  + "value is %s.",
+              storedLimitExceededBehavior, builder.limitExceededBehavior));
     }
   }
 
diff --git a/google-cloud-bigquerystorage/src/test/java/com/google/cloud/bigquery/storage/v1/it/ITBigQueryWriteManualClientTest.java b/google-cloud-bigquerystorage/src/test/java/com/google/cloud/bigquery/storage/v1/it/ITBigQueryWriteManualClientTest.java
index 36169b89b4..06ce04185a 100644
--- a/google-cloud-bigquerystorage/src/test/java/com/google/cloud/bigquery/storage/v1/it/ITBigQueryWriteManualClientTest.java
+++ b/google-cloud-bigquerystorage/src/test/java/com/google/cloud/bigquery/storage/v1/it/ITBigQueryWriteManualClientTest.java
@@ -73,6 +73,8 @@ public class ITBigQueryWriteManualClientTest {
   private static final String TABLE = "testtable";
   private static final String TABLE2 = "complicatedtable";
 
+  private static final String TEST_TRACE_ID = "DATAFLOW:job_id";
+
   private static final String DESCRIPTION = "BigQuery Write Java manual client test dataset";
 
   private static BigQueryWriteClient client;
@@ -928,6 +930,7 @@ public void testStreamWriterWithDefaultValue() throws ExecutionException, Interr
                 ProtoSchemaConverter.convert(SimpleTypeForDefaultValue.getDescriptor()))
             .setDefaultMissingValueInterpretation(MissingValueInterpretation.DEFAULT_VALUE)
             .setEnableConnectionPool(true)
+            .setTraceId(TEST_TRACE_ID)
             .build()) {
       // 1. row has both fields set.
       SimpleTypeForDefaultValue simpleTypeForDefaultValue1 =
@@ -1534,16 +1537,19 @@ public void testMultiplexingMixedLocation()
         StreamWriter.newBuilder(defaultStream1)
             .setWriterSchema(ProtoSchemaConverter.convert(FooType.getDescriptor()))
             .setEnableConnectionPool(true)
+            .setTraceId(TEST_TRACE_ID)
             .build();
     StreamWriter streamWriter2 =
         StreamWriter.newBuilder(defaultStream2)
             .setWriterSchema(ProtoSchemaConverter.convert(ComplicateType.getDescriptor()))
             .setEnableConnectionPool(true)
+            .setTraceId(TEST_TRACE_ID)
             .build();
     StreamWriter streamWriter3 =
         StreamWriter.newBuilder(defaultStream3)
             .setWriterSchema(ProtoSchemaConverter.convert(FooType.getDescriptor()))
             .setEnableConnectionPool(true)
+            .setTraceId(TEST_TRACE_ID)
             .build();
     ApiFuture<AppendRowsResponse> response1 =
         streamWriter1.append(CreateProtoRows(new String[] {"aaa"}));
@@ -1557,6 +1563,9 @@ public void testMultiplexingMixedLocation()
     assertEquals("us", streamWriter1.getLocation());
     assertEquals("us", streamWriter2.getLocation());
     assertEquals("eu", streamWriter3.getLocation());
+    streamWriter1.close();
+    streamWriter2.close();
+    streamWriter3.close();
   }
 
   @Test