chore(bigquery-jdbc): fix flaky testSetTimeout test (#13549)

Author: logachev

java-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/it/ITBase.java

@@ -33,6 +33,11 @@
 
 public class ITBase extends BigQueryJdbcBaseTest {
 
+  // This query takes 300 seconds to complete
+  public static final String query300seconds =
+      "DECLARE DELAY_TIME DATETIME; SET DELAY_TIME = DATETIME_ADD(CURRENT_DATETIME, INTERVAL 300"
+          + " SECOND); WHILE CURRENT_DATETIME < DELAY_TIME DO  END WHILE;";
+
   private static String sharedDataset;
   private static String sharedDataset2;
 

java-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/it/ITNightlyBigQueryTest.java

@@ -56,7 +56,7 @@
 import org.junit.jupiter.api.BeforeAll;
 import org.junit.jupiter.api.Test;
 
-public class ITNightlyBigQueryTest {
+public class ITNightlyBigQueryTest extends ITBase {
   static final String PROJECT_ID = ServiceOptions.getDefaultProjectId();
   static Connection bigQueryConnection;
   static Statement bigQueryStatement;
@@ -215,18 +215,13 @@ public void testQueryInterruptGracefullyStopsExplicitJob()
         DriverManager.getConnection(connection_uri + ";JobCreationMode=1", new Properties());
     Statement bigQueryStatement = bigQueryConnection.createStatement();
 
-    // This query takes 300 seconds to complete
-    String query300Seconds =
-        "DECLARE DELAY_TIME DATETIME; SET DELAY_TIME = DATETIME_ADD(CURRENT_DATETIME, INTERVAL 300"
-            + " SECOND); WHILE CURRENT_DATETIME < DELAY_TIME DO  END WHILE;";
-
     // Query will be started in the background thread & we will call cancel from current thread.
     Thread t =
         new Thread(
             () -> {
               SQLException e =
                   assertThrows(
-                      SQLException.class, () -> bigQueryStatement.execute(query300Seconds));
+                      SQLException.class, () -> bigQueryStatement.execute(ITBase.query300seconds));
               assertTrue(e.getMessage().contains("User requested cancellation"));
               threadException.set(false);
             });
@@ -254,18 +249,13 @@ public void testQueryInterruptGracefullyStopsOptionalJob()
         DriverManager.getConnection(connection_uri + ";JobCreationMode=2", new Properties());
     Statement bigQueryStatement = bigQueryConnection.createStatement();
 
-    // This query takes 300 seconds to complete
-    String query300Seconds =
-        "DECLARE DELAY_TIME DATETIME; SET DELAY_TIME = DATETIME_ADD(CURRENT_DATETIME, INTERVAL 300"
-            + " SECOND); WHILE CURRENT_DATETIME < DELAY_TIME DO  END WHILE;";
-
     // Query will be started in the background thread & we will call cancel from current thread.
     Thread t =
         new Thread(
             () -> {
               SQLException e =
                   assertThrows(
-                      SQLException.class, () -> bigQueryStatement.execute(query300Seconds));
+                      SQLException.class, () -> bigQueryStatement.execute(ITBase.query300seconds));
               assertTrue(e.getMessage().contains("Query was cancelled."));
               threadException.set(false);
             });
@@ -1684,12 +1674,4 @@ private String getSessionId() throws InterruptedException {
     Job stubJob = bigQuery.getJob(job.getJobId());
     return stubJob.getStatistics().getSessionInfo().getSessionId();
   }
-
-  private int resultSetRowCount(ResultSet resultSet) throws SQLException {
-    int rowCount = 0;
-    while (resultSet.next()) {
-      rowCount++;
-    }
-    return rowCount;
-  }
 }

java-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/it/ITStatementTest.java

@@ -16,6 +16,7 @@
 
 package com.google.cloud.bigquery.jdbc.it;
 
+import static com.google.common.truth.Truth.assertThat;
 import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.junit.jupiter.api.Assertions.assertFalse;
 import static org.junit.jupiter.api.Assertions.assertNotEquals;
@@ -48,7 +49,7 @@
 import org.junit.jupiter.api.Disabled;
 import org.junit.jupiter.api.Test;
 
-public class ITStatementTest {
+public class ITStatementTest extends ITBase {
   private static final String DEFAULT_CATALOG = ServiceOptions.getDefaultProjectId();
   private static String DATASET;
   private static Random random = new Random();
@@ -122,7 +123,7 @@ public void testExecuteQuery() throws SQLException {
     // setMaxRows Test
     statement.setMaxRows(5);
     ResultSet maxRowsResultSet = statement.executeQuery(selectQuery);
-    assertEquals(5, getSizeOfResultSet(maxRowsResultSet));
+    assertEquals(5, resultSetRowCount(maxRowsResultSet));
 
     try {
       statement.setMaxRows(0);
@@ -244,14 +245,6 @@ public void testScript() throws SQLException {
     connection.close();
   }
 
-  private int resultSetRowCount(ResultSet resultSet) throws SQLException {
-    int rowCount = 0;
-    while (resultSet.next()) {
-      rowCount++;
-    }
-    return rowCount;
-  }
-
   @Test
   public void testStringColumnLength() throws SQLException {
     String TABLE_NAME = "StringColumnLengthTable";
@@ -323,18 +316,13 @@ public void testSetTimeout() throws SQLException {
     Connection connection = DriverManager.getConnection(ITBase.connectionUrl);
     Statement statement = connection.createStatement();
 
-    String selectQuery =
-        "SELECT views FROM bigquery-public-data.wikipedia.pageviews_2020 WHERE datehour >="
-            + " '2020-01-01' LIMIT 9000000";
-
     // statement.execute(selectQuery);
     assertEquals(0, statement.getQueryTimeout());
     statement.setQueryTimeout(1);
     assertEquals(1, statement.getQueryTimeout());
-    SQLException e = assertThrows(SQLException.class, () -> statement.executeQuery(selectQuery));
-    assertEquals(
-        "BigQueryException during runQuery\nJob execution was cancelled: Job timed out",
-        e.getMessage());
+    SQLException e =
+        assertThrows(SQLException.class, () -> statement.executeQuery(ITBase.query300seconds));
+    assertThat(e.getMessage()).contains("Job execution was cancelled: Job timed out");
     statement.close();
     connection.close();
   }
@@ -388,14 +376,6 @@ public void testRangeSelectDataset() throws SQLException {
     connection.close();
   }
 
-  int getSizeOfResultSet(ResultSet resultSet) throws SQLException {
-    int count = 0;
-    while (resultSet.next()) {
-      count++;
-    }
-    return count;
-  }
-
   @Test
   public void testTemporaryDatasetLocation() throws SQLException, InterruptedException {
     String projectId = DEFAULT_CATALOG;