.

Author: yirutang

google-cloud-bigquerystorage/src/test/java/com/google/cloud/bigquery/storage/v1beta2/FakeBigQueryWrite.java

@@ -34,6 +34,10 @@ public FakeBigQueryWrite() {
  serviceImpl = new FakeBigQueryWriteImpl();
  }
 
+ public void waitForResponseScheduled() throws InterruptedException {
+ serviceImpl.waitResponseScheduled();
+ }
+
  @Override
  public List<AbstractMessage> getRequests() {
  return new LinkedList<AbstractMessage>(serviceImpl.getCapturedRequests());

google-cloud-bigquerystorage/src/test/java/com/google/cloud/bigquery/storage/v1beta2/FakeBigQueryWriteImpl.java

@@ -21,6 +21,7 @@
 import java.util.List;
 import java.util.concurrent.LinkedBlockingQueue;
 import java.util.concurrent.ScheduledExecutorService;
+import java.util.concurrent.Semaphore;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicInteger;
 import java.util.logging.Logger;
@@ -44,6 +45,7 @@ class FakeBigQueryWriteImpl extends BigQueryWriteGrpc.BigQueryWriteImplBase {
  private boolean autoPublishResponse;
  private ScheduledExecutorService executor = null;
  private Duration responseDelay = Duration.ZERO;
+ Semaphore responseCount = new Semaphore(0, true);
 
  /** Class used to save the state of a possible response. */
  private static class Response {
@@ -111,6 +113,11 @@ public void flushRows(
  }
  }
 
+ // Wait for n responses to be scheduled to sent.
+ public void waitResponseScheduled() throws InterruptedException {
+ responseCount.acquire();
+ }
+
  @Override
  public StreamObserver<AppendRowsRequest> appendRows(
  final StreamObserver<AppendRowsResponse> responseObserver) {
@@ -136,6 +143,7 @@ public void run() {
  responseDelay.toMillis(),
  TimeUnit.MILLISECONDS);
  }
+ responseCount.release();
  }
 
  @Override

google-cloud-bigquerystorage/src/test/java/com/google/cloud/bigquery/storage/v1beta2/StreamWriterTest.java

@@ -452,11 +452,11 @@ public void run() {
  assertEquals(true, t.isAlive());
  assertEquals(false, appendFuture1.isDone());
  // Wait is necessary for response to be scheduled before timer is advanced.
- Thread.sleep(5000L);
+ testBigQueryWrite.waitForResponseScheduled();
  fakeExecutor.advanceTime(Duration.ofSeconds(10));
  // The first requests gets back while the second one is blocked.
  assertEquals(2L, appendFuture1.get().getAppendResult().getOffset().getValue());
- Thread.sleep(5000L);
+ testBigQueryWrite.waitForResponseScheduled();
  // Wait is necessary for response to be scheduled before timer is advanced.
  fakeExecutor.advanceTime(Duration.ofSeconds(10));
  t.join();
@@ -496,7 +496,7 @@ public void testFlowControlBehaviorException() throws Exception {
  ApiFuture<AppendRowsResponse> appendFuture1 = sendTestMessage(writer, new String[] {"A"});
  ApiFuture<AppendRowsResponse> appendFuture2 = sendTestMessage(writer, new String[] {"B"});
  // Wait is necessary for response to be scheduled before timer is advanced.
- Thread.sleep(5000L);
+ testBigQueryWrite.waitForResponseScheduled();
  fakeExecutor.advanceTime(Duration.ofSeconds(10));
  try {
  appendFuture2.get();
@@ -944,8 +944,11 @@ public void testShutdownWithConnectionError() throws Exception {
 
  ApiFuture<AppendRowsResponse> appendFuture1 = sendTestMessage(writer, new String[] {"A"});
  ApiFuture<AppendRowsResponse> appendFuture2 = sendTestMessage(writer, new String[] {"B"});
- ApiFuture<AppendRowsResponse> appendFuture3 = sendTestMessage(writer, new String[] {"B"});
- Thread.sleep(5000L);
+ ApiFuture<AppendRowsResponse> appendFuture3 = sendTestMessage(writer, new String[] {"C"});
+ // Wait for all 3 response arrives.
+ testBigQueryWrite.waitForResponseScheduled();
+ testBigQueryWrite.waitForResponseScheduled();
+ testBigQueryWrite.waitForResponseScheduled();
  // Move the needle for responses to be sent.
  fakeExecutor.advanceTime(Duration.ofSeconds(20));
  // Shutdown writer immediately and there will be some error happened when flushing the queue.