From a69bec9c017509ef3beb975e2b6d929c12420c0c Mon Sep 17 00:00:00 2001
From: Gaole Meng <gaole@google.com>
Date: Tue, 28 Nov 2023 16:35:19 -0800
Subject: [PATCH] fix: fix the write api integration test (#2333)

* fix: an atempt to solve test failure in nightly build

* dummy PR to trigger integration test
---
 .../com/google/cloud/bigquery/storage/v1/StreamWriter.java | 6 ++++++
 .../bigquery/storage/v1/ConnectionWorkerPoolTest.java      | 7 ++++++-
 2 files changed, 12 insertions(+), 1 deletion(-)

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 e429678d8d..35e39d708f 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
@@ -586,6 +586,12 @@ ConnectionWorkerPool getTestOnlyConnectionWorkerPool() {
     return connectionWorkerPool;
   }
 
+  // A method to clear the static connectio pool to avoid making pool visible to other tests.
+  @VisibleForTesting
+  static void clearConnectionPool() {
+    connectionPoolMap.clear();
+  }
+
   /** A builder of {@link StreamWriter}s. */
   public static final class Builder {
     private static final long DEFAULT_MAX_INFLIGHT_REQUESTS = 1000L;
diff --git a/google-cloud-bigquerystorage/src/test/java/com/google/cloud/bigquery/storage/v1/ConnectionWorkerPoolTest.java b/google-cloud-bigquerystorage/src/test/java/com/google/cloud/bigquery/storage/v1/ConnectionWorkerPoolTest.java
index ccae51dcbd..5770f7734b 100644
--- a/google-cloud-bigquerystorage/src/test/java/com/google/cloud/bigquery/storage/v1/ConnectionWorkerPoolTest.java
+++ b/google-cloud-bigquerystorage/src/test/java/com/google/cloud/bigquery/storage/v1/ConnectionWorkerPoolTest.java
@@ -55,7 +55,7 @@ public class ConnectionWorkerPoolTest {
   private static MockServiceHelper serviceHelper;
   private BigQueryWriteSettings clientSettings;
 
-  private static final String TEST_TRACE_ID = "home:job1";
+  private static final String TEST_TRACE_ID = "DATAFLOW:job_id";
   private static final String TEST_STREAM_1 = "projects/p1/datasets/d1/tables/t1/streams/_default";
   private static final String TEST_STREAM_2 = "projects/p1/datasets/d1/tables/t2/streams/_default";
 
@@ -372,6 +372,7 @@ public void testToTableName() {
   @Test
   public void testCloseExternalClient()
       throws IOException, InterruptedException, ExecutionException {
+    StreamWriter.clearConnectionPool();
     // Try append 100 requests.
     long appendCount = 100L;
     // testBigQueryWrite is used to
@@ -417,6 +418,10 @@ public void testCloseExternalClient()
       assertThat(response.getAppendResult().getOffset().getValue()).isEqualTo(i);
     }
     assertThat(testBigQueryWrite.getAppendRequests().size()).isEqualTo(appendCount * 2);
+    for (int i = 0; i < streamWriterList.size(); i++) {
+      streamWriterList.get(i).close();
+    }
+    StreamWriter.clearConnectionPool();
   }
 
   private AppendRowsResponse createAppendResponse(long offset) {