From 35fe606e43a47dfe844666bf3e385f9c2e582977 Mon Sep 17 00:00:00 2001
From: Yiru Tang <yiru@google.com>
Date: Thu, 10 Dec 2020 12:30:06 -0800
Subject: [PATCH] fix: A test race in JsonStreamWriterTest  (#722)

* fix: a race condition in test

* .

* .
---
 .../storage/v1alpha2/JsonStreamWriterTest.java     | 14 ++++++--------
 .../storage/v1beta2/JsonStreamWriterTest.java      | 14 ++++++--------
 2 files changed, 12 insertions(+), 16 deletions(-)

diff --git a/google-cloud-bigquerystorage/src/test/java/com/google/cloud/bigquery/storage/v1alpha2/JsonStreamWriterTest.java b/google-cloud-bigquerystorage/src/test/java/com/google/cloud/bigquery/storage/v1alpha2/JsonStreamWriterTest.java
index 78c943e1c6..1f7ec99aa7 100644
--- a/google-cloud-bigquerystorage/src/test/java/com/google/cloud/bigquery/storage/v1alpha2/JsonStreamWriterTest.java
+++ b/google-cloud-bigquerystorage/src/test/java/com/google/cloud/bigquery/storage/v1alpha2/JsonStreamWriterTest.java
@@ -34,9 +34,7 @@
 import com.google.protobuf.Descriptors.DescriptorValidationException;
 import com.google.protobuf.Timestamp;
 import java.io.IOException;
-import java.util.Arrays;
-import java.util.HashSet;
-import java.util.UUID;
+import java.util.*;
 import java.util.concurrent.ExecutionException;
 import java.util.logging.Logger;
 import org.json.JSONArray;
@@ -777,13 +775,13 @@ public void testMultiThreadAppendNoSchemaUpdate() throws Exception {
       final JSONArray jsonArr = new JSONArray();
       jsonArr.put(foo);
 
-      final HashSet<Long> offset_sets = new HashSet<Long>();
+      final Collection<Long> offsetSets = Collections.synchronizedCollection(new HashSet<Long>());
       int thread_nums = 5;
       Thread[] thread_arr = new Thread[thread_nums];
       for (int i = 0; i < thread_nums; i++) {
         testBigQueryWrite.addResponse(
             Storage.AppendRowsResponse.newBuilder().setOffset((long) i).build());
-        offset_sets.add((long) i);
+        offsetSets.add((long) i);
         Thread t =
             new Thread(
                 new Runnable() {
@@ -792,7 +790,7 @@ public void run() {
                       ApiFuture<AppendRowsResponse> appendFuture =
                           writer.append(jsonArr, -1, /* allowUnknownFields */ false);
                       AppendRowsResponse response = appendFuture.get();
-                      offset_sets.remove(response.getOffset());
+                      offsetSets.remove(response.getOffset());
                     } catch (Exception e) {
                       LOG.severe("Thread execution failed: " + e.getMessage());
                     }
@@ -805,7 +803,7 @@ public void run() {
       for (int i = 0; i < thread_nums; i++) {
         thread_arr[i].join();
       }
-      assertTrue(offset_sets.size() == 0);
+      assertTrue(offsetSets.size() == 0);
       for (int i = 0; i < thread_nums; i++) {
         assertEquals(
             1,
@@ -842,7 +840,7 @@ public void testMultiThreadAppendWithSchemaUpdate() throws Exception {
       final JSONArray jsonArr = new JSONArray();
       jsonArr.put(foo);
 
-      final HashSet<Long> offsetSets = new HashSet<Long>();
+      final Collection<Long> offsetSets = Collections.synchronizedCollection(new HashSet<Long>());
       int numberThreads = 5;
       Thread[] thread_arr = new Thread[numberThreads];
       for (int i = 0; i < numberThreads; i++) {
diff --git a/google-cloud-bigquerystorage/src/test/java/com/google/cloud/bigquery/storage/v1beta2/JsonStreamWriterTest.java b/google-cloud-bigquerystorage/src/test/java/com/google/cloud/bigquery/storage/v1beta2/JsonStreamWriterTest.java
index 4fc3e13ef5..1c070efc79 100644
--- a/google-cloud-bigquerystorage/src/test/java/com/google/cloud/bigquery/storage/v1beta2/JsonStreamWriterTest.java
+++ b/google-cloud-bigquerystorage/src/test/java/com/google/cloud/bigquery/storage/v1beta2/JsonStreamWriterTest.java
@@ -33,9 +33,7 @@
 import com.google.protobuf.Descriptors.DescriptorValidationException;
 import com.google.protobuf.Timestamp;
 import java.io.IOException;
-import java.util.Arrays;
-import java.util.HashSet;
-import java.util.UUID;
+import java.util.*;
 import java.util.concurrent.ExecutionException;
 import java.util.logging.Logger;
 import org.json.JSONArray;
@@ -775,12 +773,12 @@ public void testMultiThreadAppendNoSchemaUpdate() throws Exception {
       final JSONArray jsonArr = new JSONArray();
       jsonArr.put(foo);
 
-      final HashSet<Long> offset_sets = new HashSet<Long>();
+      final Collection<Long> offsetSets = Collections.synchronizedCollection(new HashSet<Long>());
       int thread_nums = 5;
       Thread[] thread_arr = new Thread[thread_nums];
       for (int i = 0; i < thread_nums; i++) {
         testBigQueryWrite.addResponse(AppendRowsResponse.newBuilder().setOffset((long) i).build());
-        offset_sets.add((long) i);
+        offsetSets.add((long) i);
         Thread t =
             new Thread(
                 new Runnable() {
@@ -789,7 +787,7 @@ public void run() {
                       ApiFuture<AppendRowsResponse> appendFuture =
                           writer.append(jsonArr, -1, /* allowUnknownFields */ false);
                       AppendRowsResponse response = appendFuture.get();
-                      offset_sets.remove(response.getOffset());
+                      offsetSets.remove(response.getOffset());
                     } catch (Exception e) {
                       LOG.severe("Thread execution failed: " + e.getMessage());
                     }
@@ -802,7 +800,7 @@ public void run() {
       for (int i = 0; i < thread_nums; i++) {
         thread_arr[i].join();
       }
-      assertTrue(offset_sets.size() == 0);
+      assertTrue(offsetSets.size() == 0);
       for (int i = 0; i < thread_nums; i++) {
         assertEquals(
             1,
@@ -839,7 +837,7 @@ public void testMultiThreadAppendWithSchemaUpdate() throws Exception {
       final JSONArray jsonArr = new JSONArray();
       jsonArr.put(foo);
 
-      final HashSet<Long> offsetSets = new HashSet<Long>();
+      final Collection<Long> offsetSets = Collections.synchronizedCollection(new HashSet<Long>());
       int numberThreads = 5;
       Thread[] thread_arr = new Thread[numberThreads];
       for (int i = 0; i < numberThreads; i++) {