Fix keypair tests

snowflakedb · Aug 14, 2024 · ae91e3c · ae91e3c
1 parent e2a092d
commit ae91e3c
Showing 1 changed file with 11 additions and 2 deletions.
diff --git a/src/test/java/net/snowflake/client/jdbc/ConnectionIT.java b/src/test/java/net/snowflake/client/jdbc/ConnectionIT.java
@@ -459,8 +459,12 @@ 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));
+    }
+
+    try (Connection connection = DriverManager.getConnection(uri, properties)) {
+      assertFalse(connection.isClosed());
     } finally {
-      try (Connection connection = DriverManager.getConnection(uri, properties);
+      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));
@@ -543,7 +547,12 @@ public void testDifferentKeyLength() throws Exception {
       properties.put("privateKey", privateKey);
       try (Connection connection = DriverManager.getConnection(uri, properties);
           Statement statement = connection.createStatement()) {
-        statement.execute(String.format("alter user %s unset rsa_public_key", testUser));
+        assertFalse(connection.isClosed());
+      } finally {
+        try (Connection connection = getConnection();
+            Statement statement = connection.createStatement()) {
+          statement.execute(String.format("alter user %s unset rsa_public_key", testUser));
+        }
       }
     }
   }