From 3d74dc6519ee332723c3a1bca59387bf5afcdbfd Mon Sep 17 00:00:00 2001
From: Jelena Furundzic <jelena.furundzic@insightsoftware.com>
Date: Wed, 27 Nov 2024 22:05:53 -0800
Subject: [PATCH] fixed tests to run automatically

---
 .../client/jdbc/StatementLatestIT.java        | 60 ++++++++-----------
 1 file changed, 24 insertions(+), 36 deletions(-)

diff --git a/src/test/java/net/snowflake/client/jdbc/StatementLatestIT.java b/src/test/java/net/snowflake/client/jdbc/StatementLatestIT.java
index 12cd416c1..9b2aa377b 100644
--- 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());
   }
 }