fixed tests to run automatically

snowflakedb · Nov 28, 2024 · 3d74dc6 · 3d74dc6
1 parent 3428505
commit 3d74dc6
Showing 1 changed file with 24 additions and 36 deletions.
diff --git a/src/test/java/net/snowflake/client/jdbc/StatementLatestIT.java b/src/test/java/net/snowflake/client/jdbc/StatementLatestIT.java
@@ -299,52 +299,40 @@ public void testQueryIdIsSetOnFailedExecuteQuery() throws SQLException {
   }
 
   /**
-   * Run test manually to confirm setQueryTimeout is working for async query
+   * Test for setting query timeout on async queries. Applicable to versions after 3.20.0.
    *
    * @throws SQLException if there is an error when executing
    */
   @Test
-  @DontRunOnGithubActions
   public void testSetQueryTimeoutForAsnycQuery() throws SQLException {
-    Statement statement = connection.createStatement();
-
-    statement.unwrap(SnowflakeStatement.class).setParameter("MULTI_STATEMENT_COUNT", 0);
-    statement.setQueryTimeout(3);
-
-    String sql = "SELECT * FROM SNOWFLAKE_SAMPLE_DATA.TPCDS_SF100TCL.CUSTOMER;";
-
-    // Async
-    System.out.println("Async start time: " + java.time.LocalTime.now());
-    ResultSet resultSet = statement.unwrap(SnowflakeStatement.class).executeAsyncQuery(sql);
+    try (Statement statement = connection.createStatement()) {
+      statement.unwrap(SnowflakeStatement.class).setParameter("MULTI_STATEMENT_COUNT", 0);
+      statement.setQueryTimeout(3);
 
-    String queryID = resultSet.unwrap(SnowflakeResultSet.class).getQueryID();
-    System.out.println("Async query ID: " + queryID);
+      String sql = "SELECT * FROM SNOWFLAKE_SAMPLE_DATA.TPCDS_SF100TCL.CUSTOMER;";
 
-    QueryStatus queryStatus = QueryStatus.RUNNING;
-    while (queryStatus == QueryStatus.RUNNING || queryStatus == QueryStatus.RESUMING_WAREHOUSE) {
-      try {
-        Thread.sleep(100);
-      } catch (InterruptedException e) {
-        fail(e.getMessage());
-      }
-      queryStatus = resultSet.unwrap(SnowflakeResultSet.class).getStatus();
-      System.out.println("Current status: " + queryStatus);
-    }
+      try (ResultSet resultSet =
+          statement.unwrap(SnowflakeStatement.class).executeAsyncQuery(sql)) {
 
-    if (queryStatus == QueryStatus.SUCCESS) {
-      while (resultSet.next()) {
-        System.out.println(resultSet.getString(1));
-        System.exit(0);
-      }
-    }
+        QueryStatus queryStatus = QueryStatus.RUNNING;
+        while (queryStatus == QueryStatus.RUNNING
+            || queryStatus == QueryStatus.RESUMING_WAREHOUSE) {
+          try {
+            Thread.sleep(100);
+          } catch (InterruptedException e) {
+            fail(e.getMessage());
+          }
+          queryStatus = resultSet.unwrap(SnowflakeResultSet.class).getStatus();
+        }
 
-    if (queryStatus == QueryStatus.FAILED_WITH_ERROR) {
-      while (resultSet.next()) {
-        System.out.println(queryStatus.getErrorMessage());
-        System.exit(0);
+        if (queryStatus == QueryStatus.FAILED_WITH_ERROR) {
+          assertTrue(
+              queryStatus
+                  .getErrorMessage()
+                  .contains(
+                      "Statement reached its statement or warehouse timeout of 3 second(s) and was canceled"));
+        }
       }
     }
-
-    System.out.println("Async end time: " + java.time.LocalTime.now());
   }
 }