test: ignore if the mock server failed to return a row (#3335)

The test uses a trick in the mock server, where the mock server is
requested to freeze after returning the first row. However, when the
mock server adds the first row to the stream, it is not guaranteed to be
readable for the client, which again causes the test to hang on the
ResultSet#next() call. The gRPC libraries then execute keep-alive
requests to keep the TCP connection alive while waiting for data from
the mock server, which will never come. This caused the query to
eventually fail with a RESOURCE_EXHAUSTED error.

The tests work around this issue by just ignoring the case when the mock
server fails to return the first row, as it is something that only very
sporadically happens.

Author: olavloite

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

@@ -39,6 +39,8 @@
 import java.util.concurrent.TimeUnit;
 import java.util.stream.Collectors;
 import org.junit.After;
+import org.junit.AfterClass;
+import org.junit.BeforeClass;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 import org.junit.runners.JUnit4;
@@ -59,6 +61,16 @@ Spanner createSpanner() {
         .getService();
   }
 
+  @BeforeClass
+  public static void setWatchdogTimeout() {
+    System.setProperty("com.google.cloud.spanner.watchdogTimeoutSeconds", "1");
+  }
+
+  @AfterClass
+  public static void clearWatchdogTimeout() {
+    System.clearProperty("com.google.cloud.spanner.watchdogTimeoutSeconds");
+  }
+
   @After
   public void cleanup() {
     mockSpanner.unfreeze();
@@ -75,7 +87,13 @@ public void testBatchClient_closedSpannerWithOpenResultSet_streamsAreCancelled()
             client.batchReadOnlyTransaction(TimestampBound.strong());
         ResultSet resultSet = transaction.executeQuery(SELECT_RANDOM_STATEMENT)) {
       mockSpanner.freezeAfterReturningNumRows(1);
-      assertTrue(resultSet.next());
+      // This can sometimes fail, as the mock server may not always actually return the first row.
+      try {
+        assertTrue(resultSet.next());
+      } catch (SpannerException exception) {
+        assertEquals(ErrorCode.DEADLINE_EXCEEDED, exception.getErrorCode());
+        return;
+      }
       ((SpannerImpl) spanner).close(1, TimeUnit.MILLISECONDS);
       // This should return an error as the stream is cancelled.
       SpannerException exception = assertThrows(SpannerException.class, resultSet::next);
@@ -93,7 +111,13 @@ public void testNormalDatabaseClient_closedSpannerWithOpenResultSet_sessionsAreD
     try (ReadOnlyTransaction transaction = client.readOnlyTransaction(TimestampBound.strong());
         ResultSet resultSet = transaction.executeQuery(SELECT_RANDOM_STATEMENT)) {
       mockSpanner.freezeAfterReturningNumRows(1);
-      assertTrue(resultSet.next());
+      // This can sometimes fail, as the mock server may not always actually return the first row.
+      try {
+        assertTrue(resultSet.next());
+      } catch (SpannerException exception) {
+        assertEquals(ErrorCode.DEADLINE_EXCEEDED, exception.getErrorCode());
+        return;
+      }
       List<ExecuteSqlRequest> executeSqlRequests =
           mockSpanner.getRequestsOfType(ExecuteSqlRequest.class).stream()
               .filter(request -> request.getSql().equals(SELECT_RANDOM_STATEMENT.getSql()))