From 83e0e8cf2cdbc143c10ce0f8b053096d61d5686f Mon Sep 17 00:00:00 2001
From: Mattie Fu <mattiefu@google.com>
Date: Thu, 6 May 2021 17:03:39 -0400
Subject: [PATCH] fix: fix flaky test (#798)

---
 .../v2/stub/DynamicFlowControlCallableTest.java    | 14 +++++++++++++-
 1 file changed, 13 insertions(+), 1 deletion(-)

diff --git a/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/DynamicFlowControlCallableTest.java b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/DynamicFlowControlCallableTest.java
index 426bcdc7c..4db1f5a31 100644
--- a/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/DynamicFlowControlCallableTest.java
+++ b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/DynamicFlowControlCallableTest.java
@@ -39,6 +39,7 @@
 import java.util.Map;
 import java.util.concurrent.Future;
 import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicBoolean;
 import org.junit.Before;
 import org.junit.Rule;
 import org.junit.Test;
@@ -268,19 +269,30 @@ public ApiFuture<List<MutateRowsResponse>> futureCall(
 
   private void createFlowControlEvent(final FlowController flowController) throws Exception {
     flowController.reserve(INITIAL_ELEMENT, 0);
+    final AtomicBoolean threadStarted = new AtomicBoolean(false);
     Thread t =
         new Thread(
             new Runnable() {
               @Override
               public void run() {
                 try {
+                  threadStarted.set(true);
                   flowController.reserve(1, 0);
                 } catch (Exception e) {
                 }
               }
             });
     t.start();
-    Thread.sleep(10);
+    // Wait 5 seconds for the thread to start, and 50 milliseconds after it's started to make sure
+    // flowController.reserve(1, 0) is blocked and creates a throttling event. It should never take
+    // so long.
+    for (int i = 0; i < 1000; i++) {
+      if (threadStarted.get()) {
+        break;
+      }
+      Thread.sleep(5);
+    }
+    Thread.sleep(50);
     flowController.release(INITIAL_ELEMENT, 0);
     t.join();
     flowController.release(1, 0);