SNOW-977628: Fix flaky test ConnectionLatestIT.testIsAsyncSession (#1600

)

src/test/java/net/snowflake/client/jdbc/ConnectionLatestIT.java

@@ -18,6 +18,7 @@
 import java.nio.file.Files;
 import java.nio.file.Paths;
 import java.sql.*;
+import java.time.Duration;
 import java.util.Collections;
 import java.util.Enumeration;
 import java.util.Map;
@@ -1112,27 +1113,32 @@ public void testDownloadStreamWithFileNotFoundException() throws SQLException {
 
   @Test
   public void testIsAsyncSession() throws SQLException, InterruptedException {
+    // Run a query that takes > 4 seconds to complete
     try (Connection con = getConnection();
-        Statement statement = con.createStatement()) {
-      // Run a query that takes 5 seconds to complete
-      ResultSet rs =
-          statement
-              .unwrap(SnowflakeStatement.class)
-              .executeAsyncQuery("select count(*) from table(generator(timeLimit => 4))");
+        Statement statement = con.createStatement();
+        ResultSet rs =
+            statement
+                .unwrap(SnowflakeStatement.class)
+                .executeAsyncQuery("select count(*) from table(generator(timeLimit => 4))")) {
       // Assert that activeAsyncQueries is non-empty with running query. Session is async and not
       // safe to close
-      assertTrue(con.unwrap(SnowflakeConnectionV1.class).getSfSession().isAsyncSession());
-      assertFalse(con.unwrap(SnowflakeConnectionV1.class).getSfSession().isSafeToClose());
-      // Sleep 6 seconds to ensure query is finished running
-      TimeUnit.SECONDS.sleep(6);
+      long start = System.currentTimeMillis();
+      SnowflakeConnectionV1 snowflakeConnection = con.unwrap(SnowflakeConnectionV1.class);
+      assertTrue(snowflakeConnection.getSfSession().isAsyncSession());
+      assertFalse(snowflakeConnection.getSfSession().isSafeToClose());
+      // ensure that query is finished
+      assertTrue(rs.next());
+      // ensure query took > 4 seconds
+      assertTrue(
+          Duration.ofMillis(System.currentTimeMillis() - start).compareTo(Duration.ofSeconds(4))
+              > 0);
       // Assert that there are no longer any queries running.
       // First, assert session is safe to close. This iterates through active queries, fetches their
       // status, and removes them from the activeQueriesMap if they are no longer active.
-      assertTrue(con.unwrap(SnowflakeConnectionV1.class).getSfSession().isSafeToClose());
+      assertTrue(snowflakeConnection.getSfSession().isSafeToClose());
       // Next, assert session is no longer async (just fetches size of activeQueriesMap with no
       // other action)
-      assertFalse(con.unwrap(SnowflakeConnectionV1.class).getSfSession().isAsyncSession());
-      rs.close();
+      assertFalse(snowflakeConnection.getSfSession().isAsyncSession());
     }
   }
 }