test(spanner): add timeouts to AsyncResultSetImplTest

The AsyncResultSetImplTest seems to be hanging during some builds. This adds
timeouts to some of the test methods, and a global timeout for the tests to
prevent it from getting stuck, and instead throwing an error.

Author: olavloite

java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/AsyncResultSetImplTest.java

@@ -50,7 +50,9 @@
 import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.concurrent.atomic.AtomicInteger;
 import org.junit.Before;
+import org.junit.Rule;
 import org.junit.Test;
+import org.junit.rules.Timeout;
 import org.junit.runner.RunWith;
 import org.junit.runners.JUnit4;
 import org.mockito.Mockito;
@@ -59,6 +61,8 @@
 
 @RunWith(JUnit4.class)
 public class AsyncResultSetImplTest {
+  @Rule public final Timeout globalTimeout = Timeout.seconds(60);
+
   private ExecutorProvider mockedProvider;
   private ExecutorProvider simpleProvider;
 
@@ -198,7 +202,7 @@ public void withCallback() throws InterruptedException {
             return CallbackResponse.CONTINUE;
           });
     }
-    finishedLatch.await();
+    assertThat(finishedLatch.await(10, TimeUnit.SECONDS)).isTrue();
     // There should be between 1 and 5 callbacks, depending on the timing of the threads.
     // Normally, there should be just 1 callback.
     assertThat(callbackCounter.get()).isIn(Range.closed(1, 5));
@@ -228,7 +232,8 @@ public void callbackReceivesError() throws InterruptedException {
             return CallbackResponse.DONE;
           });
     }
-    Exception e = receivedErr.take();
+    Exception e = receivedErr.poll(10, TimeUnit.SECONDS);
+    assertThat(e).isNotNull();
     assertThat(e).isInstanceOf(SpannerException.class);
     SpannerException se = (SpannerException) e;
     assertThat(se.getErrorCode()).isEqualTo(ErrorCode.INVALID_ARGUMENT);
@@ -263,7 +268,8 @@ public void callbackReceivesErrorHalfwayThrough() throws InterruptedException {
             return CallbackResponse.DONE;
           });
     }
-    Exception e = receivedErr.take();
+    Exception e = receivedErr.poll(10, TimeUnit.SECONDS);
+    assertThat(e).isNotNull();
     assertThat(e).isInstanceOf(SpannerException.class);
     SpannerException se = (SpannerException) e;
     assertThat(se.getErrorCode()).isEqualTo(ErrorCode.INVALID_ARGUMENT);
@@ -300,8 +306,12 @@ public void pauseResume() throws InterruptedException {
             return CallbackResponse.DONE;
           });
       int rowCounter = 0;
+      long deadline = System.currentTimeMillis() + 10000; // 10 seconds
       while (!finished.get()) {
-        Object o = queue.poll(1L, TimeUnit.MILLISECONDS);
+        if (System.currentTimeMillis() > deadline) {
+          throw new RuntimeException("Test timed out waiting for finished");
+        }
+        Object o = queue.poll(10L, TimeUnit.MILLISECONDS);
         if (o != null) {
           rowCounter++;
         }
@@ -359,8 +369,12 @@ public Boolean answer(InvocationOnMock invocation) throws Throwable {
                 }
               });
       int rowCounter = 0;
+      long deadline = System.currentTimeMillis() + 10000; // 10 seconds
       while (!callbackResult.isDone()) {
-        Object o = queue.poll(1L, TimeUnit.MILLISECONDS);
+        if (System.currentTimeMillis() > deadline) {
+          throw new RuntimeException("Test timed out waiting for callbackResult");
+        }
+        Object o = queue.poll(10L, TimeUnit.MILLISECONDS);
         if (o != null) {
           rowCounter++;
         }
@@ -453,8 +467,12 @@ public void cancel() throws InterruptedException {
             return CallbackResponse.DONE;
           });
       int rowCounter = 0;
+      long deadline = System.currentTimeMillis() + 10000; // 10 seconds
       while (!finished.get()) {
-        Object o = queue.poll(1L, TimeUnit.MILLISECONDS);
+        if (System.currentTimeMillis() > deadline) {
+          throw new RuntimeException("Test timed out waiting for finished");
+        }
+        Object o = queue.poll(10L, TimeUnit.MILLISECONDS);
         if (o != null) {
           rowCounter++;
         }