chore(bigquery-jdbc): fix flaky testSetTimeout test

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
+  protected 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(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(query300seconds));
               assertTrue(e.getMessage().contains("Query was cancelled."));
               threadException.set(false);
             });

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

@@ -48,7 +48,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();
@@ -323,15 +323,12 @@ 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));
+    SQLException e =
+        assertThrows(SQLException.class, () -> statement.executeQuery(query300seconds));
     assertEquals(
         "BigQueryException during runQuery\nJob execution was cancelled: Job timed out",
         e.getMessage());
@@ -388,14 +385,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;