SNOW-806332 Fix test timeouts (#1380)

* SNOW-806332 Fix test timeouts

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

@@ -216,7 +216,7 @@ public void testExecuteLargeBatch() throws SQLException {
 
   @Test
   public void testRemoveExtraDescribeCalls() throws SQLException {
-    Connection connection = getConnection();
+    Connection connection = init();
     Statement statement = connection.createStatement();
     statement.execute("create or replace table test_uuid_with_bind(c1 number)");
 
@@ -259,7 +259,7 @@ public void testRemoveExtraDescribeCalls() throws SQLException {
 
   @Test
   public void testRemoveExtraDescribeCallsSanityCheck() throws SQLException {
-    Connection connection = getConnection();
+    Connection connection = init();
     PreparedStatement preparedStatement =
         connection.prepareStatement(
             "create or replace table test_uuid_with_bind(c1 number, c2 string)");
@@ -297,7 +297,7 @@ public void testRemoveExtraDescribeCallsSanityCheck() throws SQLException {
 
   @Test
   public void testAlreadyDescribedMultipleResults() throws SQLException {
-    Connection connection = getConnection();
+    Connection connection = init();
     PreparedStatement prepStatement = connection.prepareStatement(insertSQL);
     bindOneParamSet(prepStatement, 1, 1.22222, (float) 1.2, "test", 12121212121L, (short) 12);
     prepStatement.execute();
@@ -317,30 +317,4 @@ public void testAlreadyDescribedMultipleResults() throws SQLException {
     assertTrue(rs.next());
     assertTrue(prepStatement.unwrap(SnowflakePreparedStatementV1.class).isAlreadyDescribed());
   }
-
-  @Test
-  public void testExecuteLargeBatchOverMaxInt() throws SQLException {
-    Connection connection = null;
-    try {
-      connection = getConnection();
-      try (Statement statement = connection.createStatement()) {
-        statement.execute("create or replace table over_int(id decimal);");
-        for (int i = 1; i <= 11; i++) {
-          statement.execute(
-              "insert into over_int (SELECT seq8() FROM table(generator(rowCount => 200000000)));");
-        }
-        String sqlStatement = "UPDATE over_int SET id = 200";
-
-        PreparedStatement pstmt = connection.prepareStatement(sqlStatement);
-        pstmt.addBatch();
-        long[] queryResult = pstmt.executeLargeBatch();
-        assertEquals(1, queryResult.length);
-      }
-    } finally {
-      if (connection != null) {
-        connection.createStatement().execute("drop table if exists over_int ");
-        connection.close();
-      }
-    }
-  }
 }

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

@@ -0,0 +1,85 @@
+/*
+ * Copyright (c) 2012-2019 Snowflake Computing Inc. All rights reserved.
+ */
+package net.snowflake.client.jdbc;
+
+import static org.junit.Assert.assertEquals;
+import static org.mockito.Mockito.spy;
+
+import java.sql.Connection;
+import java.sql.PreparedStatement;
+import java.sql.SQLException;
+import java.util.Map;
+import net.snowflake.client.ConditionalIgnoreRule;
+import net.snowflake.client.RunningOnGithubAction;
+import net.snowflake.client.category.TestCategoryStatement;
+import net.snowflake.client.core.ExecTimeTelemetryData;
+import org.junit.Test;
+import org.junit.experimental.categories.Category;
+import org.mockito.Mockito;
+
+@Category(TestCategoryStatement.class)
+public class PreparedStatementLargeUpdateLatestIT extends BaseJDBCTest {
+
+  /**
+   * Test that large update larger than MAX_INT returns correctly
+   *
+   * @throws Throwable
+   */
+  @Test
+  @ConditionalIgnoreRule.ConditionalIgnore(condition = RunningOnGithubAction.class)
+  public void testLargeUpdate() throws Throwable {
+    try (Connection con = getConnection()) {
+      long expectedUpdateRows = (long) Integer.MAX_VALUE + 10L;
+      con.createStatement().execute("create or replace table  test_large_update(c1 boolean)");
+      PreparedStatement st =
+          con.prepareStatement(
+              "insert into test_large_update select true from table(generator(rowcount=>"
+                  + expectedUpdateRows
+                  + "))");
+      PreparedStatement spyp = spy(st);
+      // Mock internal method which returns rowcount
+      Mockito.doReturn(expectedUpdateRows)
+          .when((SnowflakePreparedStatementV1) spyp)
+          .executeUpdateInternal(
+              Mockito.any(String.class),
+              Mockito.any(Map.class),
+              Mockito.any(boolean.class),
+              Mockito.any(ExecTimeTelemetryData.class));
+      long updatedRows = spyp.executeLargeUpdate();
+      assertEquals(expectedUpdateRows, updatedRows);
+      con.createStatement().execute("drop table if exists test_large_update");
+    }
+  }
+
+  /**
+   * Test that when a batch size larger than MAX_INT is returned, the result is returned as a long.
+   *
+   * @throws SQLException
+   */
+  @Test
+  @ConditionalIgnoreRule.ConditionalIgnore(condition = RunningOnGithubAction.class)
+  public void testExecuteLargeBatchOverIntMax() throws SQLException {
+    try (Connection connection = getConnection()) {
+      connection
+          .createStatement()
+          .execute("create or replace table over_int_table (val string, id int)");
+      PreparedStatement pstmt = connection.prepareStatement("UPDATE over_int_table SET ID=200");
+      PreparedStatement spyp = spy(pstmt);
+      long numRows = Integer.MAX_VALUE + 10L;
+      // Mock internal method which returns rowcount
+      Mockito.doReturn(numRows)
+          .when((SnowflakePreparedStatementV1) spyp)
+          .executeUpdateInternal(
+              Mockito.any(String.class),
+              Mockito.any(Map.class),
+              Mockito.any(boolean.class),
+              Mockito.any(ExecTimeTelemetryData.class));
+      pstmt.addBatch();
+      long[] queryResult = spyp.executeLargeBatch();
+      assertEquals(1, queryResult.length);
+      assertEquals(numRows, queryResult[0]);
+      connection.createStatement().execute("drop table if exists over_int_table");
+    }
+  }
+}