Fix keypair tests

snowflakedb · Aug 20, 2024 · b26cec6 · b26cec6
1 parent e2a092d
commit b26cec6
Showing 1 changed file with 17 additions and 8 deletions.
diff --git a/src/test/java/net/snowflake/client/jdbc/ConnectionIT.java b/src/test/java/net/snowflake/client/jdbc/ConnectionIT.java
@@ -459,13 +459,17 @@ public void testConnectUsingKeyPair() throws Exception {
       String encodePublicKey2 = Base64.encodeBase64String(publicKey2.getEncoded());
       statement.execute(
           String.format("alter user %s set rsa_public_key_2='%s'", testUser, encodePublicKey2));
-    } finally {
-      try (Connection connection = DriverManager.getConnection(uri, properties);
-          Statement statement = connection.createStatement()) {
-        statement.execute("use role accountadmin");
-        statement.execute(String.format("alter user %s unset rsa_public_key", testUser));
-        statement.execute(String.format("alter user %s unset rsa_public_key_2", testUser));
-      }
+    }
+
+    try (Connection connection = DriverManager.getConnection(uri, properties)) {
+      assertFalse(connection.isClosed());
+    }
+
+    try (Connection connection = getConnection();
+        Statement statement = connection.createStatement()) {
+      statement.execute("use role accountadmin");
+      statement.execute(String.format("alter user %s unset rsa_public_key", testUser));
+      statement.execute(String.format("alter user %s unset rsa_public_key_2", testUser));
     }
   }
 
@@ -541,8 +545,13 @@ public void testDifferentKeyLength() throws Exception {
 
       // test correct private key one
       properties.put("privateKey", privateKey);
-      try (Connection connection = DriverManager.getConnection(uri, properties);
+      try (Connection connection = DriverManager.getConnection(uri, properties)) {
+        assertFalse(connection.isClosed());
+      }
+
+      try (Connection connection = getConnection();
           Statement statement = connection.createStatement()) {
+        statement.execute("use role accountadmin");
         statement.execute(String.format("alter user %s unset rsa_public_key", testUser));
       }
     }