Add missing test object cleanup calls

src/test/java/com/microsoft/sqlserver/jdbc/AlwaysEncrypted/BulkCopySendTemporalDataTypesAsStringAETest.java

@@ -167,7 +167,9 @@ public static void setupTest() throws SQLException {
 
     @AfterAll
     public static void cleanTest() throws SQLException {
-        try (Connection con = getConnection(); Statement stmt = con.createStatement()) {
+        // Make sure to clean objects against the AE connection string, not the common connection
+        try (Connection con = PrepUtil.getConnection(AETestConnectionString, AEInfo);
+                Statement stmt = con.createStatement()) {
             TestUtils.dropTableIfExists(destTableNameAE, stmt);
         }
     }

src/test/java/com/microsoft/sqlserver/jdbc/AlwaysEncrypted/CallableStatementTest.java

@@ -9,10 +9,10 @@
 import static org.junit.jupiter.api.Assertions.fail;
 
 import java.math.BigDecimal;
+import java.sql.CallableStatement;
 import java.sql.Connection;
 import java.sql.Date;
 import java.sql.PreparedStatement;
-import java.sql.CallableStatement;
 import java.sql.ResultSet;
 import java.sql.SQLException;
 import java.sql.Statement;
@@ -160,8 +160,9 @@ public static void initValues() throws Exception {
 
     @AfterAll
     public static void dropAll() throws Exception {
-        dropTables();
+        // Drop procedures before tables or table drop will fail
         dropProcedures();
+        dropTables();
     }
 
     @ParameterizedTest

src/test/java/com/microsoft/sqlserver/jdbc/AlwaysEncrypted/EnclaveTest.java

@@ -4,8 +4,8 @@
  */
 package com.microsoft.sqlserver.jdbc.AlwaysEncrypted;
 
-import static org.junit.jupiter.api.Assertions.assertTrue;
 import static org.junit.Assert.fail;
+import static org.junit.jupiter.api.Assertions.assertTrue;
 
 import java.sql.Connection;
 import java.sql.DriverManager;
@@ -347,8 +347,12 @@ public void testChar(String serverName, String url, String protocol) throws Exce
                 SQLServerStatement stmt = (SQLServerStatement) con.createStatement()) {
             TestUtils.dropTableIfExists(CHAR_TABLE_AE, stmt);
             createTable(CHAR_TABLE_AE, cekJks, charTable);
-            populateCharNormalCase(createCharValues(false));
-            testAlterColumnEncryption(stmt, CHAR_TABLE_AE, charTable, cekJks);
+            try {
+                populateCharNormalCase(createCharValues(false));
+                testAlterColumnEncryption(stmt, CHAR_TABLE_AE, charTable, cekJks);
+            } finally {
+                TestUtils.dropTableIfExists(CHAR_TABLE_AE, stmt);
+            }
         }
     }
 
@@ -364,8 +368,12 @@ public void testCharAkv(String serverName, String url, String protocol) throws E
                 SQLServerStatement stmt = (SQLServerStatement) con.createStatement()) {
             TestUtils.dropTableIfExists(CHAR_TABLE_AE, stmt);
             createTable(CHAR_TABLE_AE, cekAkv, charTable);
-            populateCharNormalCase(createCharValues(false));
-            testAlterColumnEncryption(stmt, CHAR_TABLE_AE, charTable, cekAkv);
+            try {
+                populateCharNormalCase(createCharValues(false));
+                testAlterColumnEncryption(stmt, CHAR_TABLE_AE, charTable, cekAkv);
+            } finally {
+                TestUtils.dropTableIfExists(CHAR_TABLE_AE, connection.createStatement());
+            }
         }
     }
 
@@ -387,6 +395,8 @@ public void testAEFMTOnly(String serverName, String url, String protocol) throws
             try (PreparedStatement p = c.prepareStatement(sql)) {
                 ParameterMetaData pmd = p.getParameterMetaData();
                 assertTrue(48 == pmd.getParameterCount(), "parameter count: " + pmd.getParameterCount());
+            } finally {
+                TestUtils.dropTableIfExists(NUMERIC_TABLE_AE, s);
             }
         }
     }
@@ -401,14 +411,18 @@ public void testAlter(String serverName, String url, String protocol) throws Exc
         try (SQLServerConnection c = PrepUtil.getConnection(AETestConnectionString, AEInfo);
                 Statement s = c.createStatement()) {
             createTable(CHAR_TABLE_AE, cekJks, varcharTableSimple);
-            PreparedStatement pstmt = c.prepareStatement("INSERT INTO " + CHAR_TABLE_AE + " VALUES (?,?,?)");
-            pstmt.setString(1, "a");
-            pstmt.setString(2, "b");
-            pstmt.setString(3, "test");
-            pstmt.execute();
-            pstmt = c.prepareStatement("ALTER TABLE " + CHAR_TABLE_AE
-                    + " ALTER COLUMN RandomizedVarchar VARCHAR(20) NULL WITH (ONLINE = ON)");
-            pstmt.execute();
+            try {
+                PreparedStatement pstmt = c.prepareStatement("INSERT INTO " + CHAR_TABLE_AE + " VALUES (?,?,?)");
+                pstmt.setString(1, "a");
+                pstmt.setString(2, "b");
+                pstmt.setString(3, "test");
+                pstmt.execute();
+                pstmt = c.prepareStatement("ALTER TABLE " + CHAR_TABLE_AE
+                        + " ALTER COLUMN RandomizedVarchar VARCHAR(20) NULL WITH (ONLINE = ON)");
+                pstmt.execute();
+            } finally {
+                TestUtils.dropTableIfExists(CHAR_TABLE_AE, s);
+            }
         }
     }
 
@@ -422,19 +436,23 @@ public void testNumericRichQuery(String serverName, String url, String protocol)
         try (SQLServerConnection c = PrepUtil.getConnection(AETestConnectionString, AEInfo);
                 Statement s = c.createStatement()) {
             createTable(NUMERIC_TABLE_AE, cekJks, numericTableSimple);
-            PreparedStatement pstmt = c.prepareStatement("INSERT INTO " + NUMERIC_TABLE_AE + " VALUES (?,?,?)");
-            pstmt.setInt(1, 1);
-            pstmt.setInt(2, 2);
-            pstmt.setInt(3, 3);
-            pstmt.execute();
-            pstmt = c.prepareStatement("SELECT * FROM " + NUMERIC_TABLE_AE + " WHERE RANDOMIZEDInt = ?");
-            pstmt.setInt(1, 3);
-            try (ResultSet rs = pstmt.executeQuery()) {
-                while (rs.next()) {
-                    assertTrue(1 == rs.getInt(1), "rs.getInt(1)=" + rs.getInt(1));
-                    assertTrue(2 == rs.getInt(2), "rs.getInt(2)=" + rs.getInt(2));
-                    assertTrue(3 == rs.getInt(3), "rs.getInt(3)=" + rs.getInt(3));
+            try {
+                PreparedStatement pstmt = c.prepareStatement("INSERT INTO " + NUMERIC_TABLE_AE + " VALUES (?,?,?)");
+                pstmt.setInt(1, 1);
+                pstmt.setInt(2, 2);
+                pstmt.setInt(3, 3);
+                pstmt.execute();
+                pstmt = c.prepareStatement("SELECT * FROM " + NUMERIC_TABLE_AE + " WHERE RANDOMIZEDInt = ?");
+                pstmt.setInt(1, 3);
+                try (ResultSet rs = pstmt.executeQuery()) {
+                    while (rs.next()) {
+                        assertTrue(1 == rs.getInt(1), "rs.getInt(1)=" + rs.getInt(1));
+                        assertTrue(2 == rs.getInt(2), "rs.getInt(2)=" + rs.getInt(2));
+                        assertTrue(3 == rs.getInt(3), "rs.getInt(3)=" + rs.getInt(3));
+                    }
                 }
+            } finally {
+                TestUtils.dropTableIfExists(NUMERIC_TABLE_AE, s);
             }
         }
     }
@@ -449,20 +467,23 @@ public void testStringRichQuery(String serverName, String url, String protocol)
         try (SQLServerConnection c = PrepUtil.getConnection(AETestConnectionString, AEInfo);
                 Statement s = c.createStatement()) {
             createTable(CHAR_TABLE_AE, cekJks, varcharTableSimple);
-
-            PreparedStatement pstmt = c.prepareStatement("INSERT INTO " + CHAR_TABLE_AE + " VALUES (?,?,?)");
-            pstmt.setString(1, "a");
-            pstmt.setString(2, "b");
-            pstmt.setString(3, "test");
-            pstmt.execute();
-            pstmt = c.prepareStatement("SELECT * FROM " + CHAR_TABLE_AE + " WHERE RANDOMIZEDVarchar LIKE ?");
-            pstmt.setString(1, "t%");
-            try (ResultSet rs = pstmt.executeQuery()) {
-                while (rs.next()) {
-                    assertTrue(rs.getString(1).equalsIgnoreCase("a"), "rs.getString(1)=" + rs.getString(1));
-                    assertTrue(rs.getString(2).equalsIgnoreCase("b"), "rs.getString(2)=" + rs.getString(2));
-                    assertTrue(rs.getString(3).equalsIgnoreCase("test"), "rs.getString(3)=" + rs.getString(3));
+            try {
+                PreparedStatement pstmt = c.prepareStatement("INSERT INTO " + CHAR_TABLE_AE + " VALUES (?,?,?)");
+                pstmt.setString(1, "a");
+                pstmt.setString(2, "b");
+                pstmt.setString(3, "test");
+                pstmt.execute();
+                pstmt = c.prepareStatement("SELECT * FROM " + CHAR_TABLE_AE + " WHERE RANDOMIZEDVarchar LIKE ?");
+                pstmt.setString(1, "t%");
+                try (ResultSet rs = pstmt.executeQuery()) {
+                    while (rs.next()) {
+                        assertTrue(rs.getString(1).equalsIgnoreCase("a"), "rs.getString(1)=" + rs.getString(1));
+                        assertTrue(rs.getString(2).equalsIgnoreCase("b"), "rs.getString(2)=" + rs.getString(2));
+                        assertTrue(rs.getString(3).equalsIgnoreCase("test"), "rs.getString(3)=" + rs.getString(3));
+                    }
                 }
+            } finally {
+                TestUtils.dropTableIfExists(CHAR_TABLE_AE, s);
             }
         }
     }
@@ -477,27 +498,36 @@ public void testAlterNoEncrypt(String serverName, String url, String protocol) t
         try (SQLServerConnection c = PrepUtil.getConnection(AETestConnectionString, AEInfo);
                 Statement s = c.createStatement()) {
             createTable(CHAR_TABLE_AE, cekJks, varcharTableSimple);
-            PreparedStatement pstmt = c.prepareStatement("INSERT INTO " + CHAR_TABLE_AE + " VALUES (?,?,?)");
-            pstmt.setString(1, "a");
-            pstmt.setString(2, "b");
-            pstmt.setString(3, "test");
-            pstmt.execute();
-        }
-        String testConnectionString = TestUtils.removeProperty(AETestConnectionString,
-                Constants.ENCLAVE_ATTESTATIONURL);
-        testConnectionString = TestUtils.removeProperty(testConnectionString, Constants.ENCLAVE_ATTESTATIONPROTOCOL);
-        try (Connection c = DriverManager.getConnection(testConnectionString)) {
-            PreparedStatement pstmt = c.prepareStatement("ALTER TABLE " + CHAR_TABLE_AE
-                    + " ALTER COLUMN RandomizedVarchar VARCHAR(20) NULL WITH (ONLINE = ON)");
-            pstmt.execute();
-        } catch (SQLException e) {
-            assertTrue(e.getMessage().contains(TestResource.getResource("R_enclaveNotEnabled")), e.getMessage());
+            try {
+                PreparedStatement pstmt = c.prepareStatement("INSERT INTO " + CHAR_TABLE_AE + " VALUES (?,?,?)");
+                pstmt.setString(1, "a");
+                pstmt.setString(2, "b");
+                pstmt.setString(3, "test");
+                pstmt.execute();
+
+                String testConnectionString = TestUtils.removeProperty(AETestConnectionString,
+                        Constants.ENCLAVE_ATTESTATIONURL);
+                testConnectionString = TestUtils.removeProperty(testConnectionString,
+                        Constants.ENCLAVE_ATTESTATIONPROTOCOL);
+                try (Connection c1 = DriverManager.getConnection(testConnectionString)) {
+                    PreparedStatement pstmt1 = c1.prepareStatement("ALTER TABLE " + CHAR_TABLE_AE
+                            + " ALTER COLUMN RandomizedVarchar VARCHAR(20) NULL WITH (ONLINE = ON)");
+                    pstmt1.execute();
+                } catch (SQLException e) {
+                    assertTrue(e.getMessage().contains(TestResource.getResource("R_enclaveNotEnabled")),
+                            e.getMessage());
+                }
+            } finally {
+                TestUtils.dropTableIfExists(CHAR_TABLE_AE, s);
+            }
         }
     }
 
     @AfterAll
     public static void dropAll() throws Exception {
-        try (Statement stmt = connection.createStatement()) {
+        // Make sure to clean objects against the AE connection string, not the common connection
+        try (SQLServerConnection cn = PrepUtil.getConnection(AETestConnectionString, AEInfo);
+                Statement stmt = cn.createStatement()) {
             TestUtils.dropTableIfExists(CHAR_TABLE_AE, stmt);
             TestUtils.dropTableIfExists(NUMERIC_TABLE_AE, stmt);
             TestUtils.dropTableIfExists(BINARY_TABLE_AE, stmt);

src/test/java/com/microsoft/sqlserver/jdbc/AlwaysEncrypted/MSITest.java

@@ -16,16 +16,15 @@
 import java.util.Map;
 import java.util.Properties;
 
-import com.azure.identity.CredentialUnavailableException;
-import com.azure.identity.ManagedIdentityCredential;
-import com.azure.identity.ManagedIdentityCredentialBuilder;
-
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Tag;
 import org.junit.jupiter.api.Test;
 import org.junit.platform.runner.JUnitPlatform;
 import org.junit.runner.RunWith;
 
+import com.azure.identity.CredentialUnavailableException;
+import com.azure.identity.ManagedIdentityCredential;
+import com.azure.identity.ManagedIdentityCredentialBuilder;
 import com.microsoft.sqlserver.jdbc.SQLServerColumnEncryptionAzureKeyVaultProvider;
 import com.microsoft.sqlserver.jdbc.SQLServerColumnEncryptionKeyStoreProvider;
 import com.microsoft.sqlserver.jdbc.SQLServerConnection;
@@ -454,6 +453,8 @@ private void testNumericAKV(String connStr) throws SQLException {
                     AECommon.testGetBigDecimal(rs, numberOfColumns, values);
                     AECommon.testWithSpecifiedtype(rs, numberOfColumns, values);
                 }
+            } finally {
+                TestUtils.dropTableIfExists(NUMERIC_TABLE_AE, stmt);
             }
         }
     }

src/test/java/com/microsoft/sqlserver/jdbc/AlwaysEncrypted/RegressionAlwaysEncryptedTest.java

@@ -311,5 +311,6 @@ public static void dropTables(Statement stmt) throws SQLException {
         TestUtils.dropTableIfExists(DATE_TABLE_AE, stmt);
         TestUtils.dropTableIfExists(CHAR_TABLE_AE, stmt);
         TestUtils.dropTableIfExists(NUMERIC_TABLE_AE, stmt);
+        TestUtils.dropTableIfExists(VARY_STRING_TABLE_AE, stmt);
     }
 }

src/test/java/com/microsoft/sqlserver/jdbc/TestUtils.java

@@ -507,8 +507,10 @@ public static void dropTypeIfExists(String typeName, java.sql.Statement stmt) th
      * @throws SQLException
      */
     public static void dropUserDefinedTypeIfExists(String typeName, Statement stmt) throws SQLException {
-        stmt.executeUpdate("IF EXISTS (select * from sys.types where name = '" + escapeSingleQuotes(typeName)
-                + "') DROP TYPE " + typeName);
+        String input = AbstractSQLGenerator.unEscapeIdentifier(typeName);
+        String sql = "IF EXISTS (select * from sys.types where name = '" + escapeSingleQuotes(input) + "') DROP TYPE "
+                + AbstractSQLGenerator.escapeIdentifier(input);
+        stmt.executeUpdate(sql);
     }
 
     /**