Updates for running tests with managed identity (#2416)

microsoft · May 22, 2024 · 82bd335 · 82bd335
1 parent 7deb06f
commit 82bd335
Show file tree

Hide file tree

Showing 14 changed files with 231 additions and 71 deletions.
diff --git a/src/test/java/com/microsoft/sqlserver/jdbc/AlwaysEncrypted/AESetup.java b/src/test/java/com/microsoft/sqlserver/jdbc/AlwaysEncrypted/AESetup.java
@@ -332,7 +332,7 @@ protected static void createTable(String tableName, String cekName, String table
             TestUtils.dropTableIfExists(tableName, stmt);
             sql = String.format(createSql, tableName, sql);
             stmt.execute(sql);
-            stmt.execute("DBCC FREEPROCCACHE");
+            TestUtils.freeProcCache(stmt);
         } catch (SQLException e) {
             fail(e.getMessage());
         }
@@ -366,7 +366,7 @@ protected static void createPrecisionTable(String tableName, String table[][], S
             }
             sql = String.format(createSql, tableName, sql);
             stmt.execute(sql);
-            stmt.execute("DBCC FREEPROCCACHE");
+            TestUtils.freeProcCache(stmt);
         } catch (SQLException e) {
             fail(e.getMessage());
         }
@@ -394,7 +394,7 @@ protected static void createScaleTable(String tableName, String table[][], Strin
 
             sql = String.format(createSql, tableName, sql);
             stmt.execute(sql);
-            stmt.execute("DBCC FREEPROCCACHE");
+            TestUtils.freeProcCache(stmt);
         } catch (SQLException e) {
             fail(e.getMessage());
         }

diff --git a/src/test/java/com/microsoft/sqlserver/jdbc/AlwaysEncrypted/CallableStatementTest.java b/src/test/java/com/microsoft/sqlserver/jdbc/AlwaysEncrypted/CallableStatementTest.java
@@ -2201,7 +2201,7 @@ protected static void createDateTableCallableStatement(String cekName) throws SQ
                 SQLServerStatement stmt = (SQLServerStatement) con.createStatement()) {
             TestUtils.dropTableIfExists(DATE_TABLE_AE, stmt);
             stmt.execute(sql);
-            stmt.execute("DBCC FREEPROCCACHE");
+            TestUtils.freeProcCache(stmt);
         } catch (SQLException e) {
             fail(e.getMessage());
         }

diff --git a/src/test/java/com/microsoft/sqlserver/jdbc/SQLServerConnectionTest.java b/src/test/java/com/microsoft/sqlserver/jdbc/SQLServerConnectionTest.java
@@ -481,7 +481,13 @@ public void testConnectCountInLoginAndCorrectRetryCount() {
                 assertTrue(con == null, TestResource.getResource("R_shouldNotConnect"));
             }
         } catch (Exception e) {
-            assertTrue(e.getMessage().contains(TestResource.getResource("R_cannotOpenDatabase")), e.getMessage());
+            assertTrue(
+                    e.getMessage().contains(TestResource.getResource("R_cannotOpenDatabase"))
+                            || (TestUtils.getProperty(connectionString, "msiClientId") != null && (e.getMessage()
+                                    .toLowerCase().contains(TestResource.getResource("R_loginFailedMI").toLowerCase())
+                                    || e.getMessage().toLowerCase()
+                                            .contains(TestResource.getResource("R_MInotAvailable").toLowerCase()))),
+                    e.getMessage());
             long totalTime = System.currentTimeMillis() - timerStart;
 
             // Maximum is unknown, but is needs to be less than longLoginTimeout or else this is an issue.
@@ -756,13 +762,22 @@ public void testIncorrectDatabase() throws SQLException {
                 assertTrue(timeDiff <= milsecs, form.format(msgArgs));
             }
         } catch (Exception e) {
-            assertTrue(e.getMessage().contains(TestResource.getResource("R_cannotOpenDatabase")), e.getMessage());
+            assertTrue(
+                    e.getMessage().contains(TestResource.getResource("R_cannotOpenDatabase"))
+                            || (TestUtils.getProperty(connectionString, "msiClientId") != null
+                                    && e.getMessage().toLowerCase()
+                                            .contains(TestResource.getResource("R_loginFailedMI").toLowerCase())),
+                    e.getMessage());
             timerEnd = System.currentTimeMillis();
         }
     }
 
     @Test
     public void testIncorrectUserName() throws SQLException {
+        String auth = TestUtils.getProperty(connectionString, "authentication");
+        org.junit.Assume.assumeTrue(auth != null
+                && (auth.equalsIgnoreCase("SqlPassword") || auth.equalsIgnoreCase("ActiveDirectoryPassword")));
+
         long timerStart = 0;
         long timerEnd = 0;
         final long milsecs = threshHoldForNoRetryInMilliseconds;
@@ -780,13 +795,22 @@ public void testIncorrectUserName() throws SQLException {
                 assertTrue(timeDiff <= milsecs, form.format(msgArgs));
             }
         } catch (Exception e) {
-            assertTrue(e.getMessage().contains(TestResource.getResource("R_loginFailed")));
+            assertTrue(
+                    e.getMessage().contains(TestResource.getResource("R_loginFailed"))
+                            || (TestUtils.getProperty(connectionString, "msiClientId") != null
+                                    && e.getMessage().toLowerCase()
+                                            .contains(TestResource.getResource("R_loginFailedMI").toLowerCase())),
+                    e.getMessage());
             timerEnd = System.currentTimeMillis();
         }
     }
 
     @Test
     public void testIncorrectPassword() throws SQLException {
+        String auth = TestUtils.getProperty(connectionString, "authentication");
+        org.junit.Assume.assumeTrue(auth != null
+                && (auth.equalsIgnoreCase("SqlPassword") || auth.equalsIgnoreCase("ActiveDirectoryPassword")));
+
         long timerStart = 0;
         long timerEnd = 0;
         final long milsecs = threshHoldForNoRetryInMilliseconds;
@@ -804,7 +828,12 @@ public void testIncorrectPassword() throws SQLException {
                 assertTrue(timeDiff <= milsecs, form.format(msgArgs));
             }
         } catch (Exception e) {
-            assertTrue(e.getMessage().contains(TestResource.getResource("R_loginFailed")));
+            assertTrue(
+                    e.getMessage().contains(TestResource.getResource("R_loginFailed"))
+                            || (TestUtils.getProperty(connectionString, "msiClientId") != null
+                                    && e.getMessage().toLowerCase()
+                                            .contains(TestResource.getResource("R_loginFailedMI").toLowerCase())),
+                    e.getMessage());
             timerEnd = System.currentTimeMillis();
         }
     }

diff --git a/src/test/java/com/microsoft/sqlserver/jdbc/TestResource.java b/src/test/java/com/microsoft/sqlserver/jdbc/TestResource.java
@@ -212,5 +212,9 @@ protected Object[][] getContents() {
             {"R_failedFedauth", "Failed to acquire fedauth token: "},
             {"R_noLoginModulesConfiguredForJdbcDriver",
                     "javax.security.auth.login.LoginException (No LoginModules configured for SQLJDBCDriver)"},
-            {"R_unexpectedThreadCount", "Thread count is higher than expected."}};
+            {"R_unexpectedThreadCount", "Thread count is higher than expected."},
+            {"R_expectedClassDoesNotMatchActualClass",
+                    "Expected column class {0} does not match actual column class {1} for column {2}."},
+            {"R_loginFailedMI", "Login failed for user '<token-identified principal>'"},
+            {"R_MInotAvailable", "Managed Identity authentication is not available"},};
 }
diff --git a/src/test/java/com/microsoft/sqlserver/jdbc/TestUtils.java b/src/test/java/com/microsoft/sqlserver/jdbc/TestUtils.java
@@ -523,7 +523,31 @@ public static void dropDatabaseIfExists(String databaseName, String connectionSt
      */
     public static void dropSchemaIfExists(String schemaName, Statement stmt) throws SQLException {
         stmt.execute("if EXISTS (SELECT * FROM sys.schemas where name = '" + escapeSingleQuotes(schemaName)
-                + "') drop schema " + AbstractSQLGenerator.escapeIdentifier(schemaName));
+                + "') DROP SCHEMA" + AbstractSQLGenerator.escapeIdentifier(schemaName));
+    }
+
+    /**
+     * mimic "DROP USER..."
+     * 
+     * @param userName
+     * @param stmt
+     * @throws SQLException
+     */
+    public static void dropUserIfExists(String userName, Statement stmt) throws SQLException {
+        stmt.execute("IF EXISTS (SELECT * FROM sys.sysusers where name = '" + escapeSingleQuotes(userName)
+                + "') DROP USER " + AbstractSQLGenerator.escapeIdentifier(userName));
+    }
+
+    /**
+     * mimic "DROP LOGIN..."
+     * 
+     * @param userName
+     * @param stmt
+     * @throws SQLException
+     */
+    public static void dropLoginIfExists(String userName, Statement stmt) throws SQLException {
+        stmt.execute("IF EXISTS (SELECT * FROM sys.sysusers where name = '" + escapeSingleQuotes(userName)
+                + "') DROP LOGIN " + AbstractSQLGenerator.escapeIdentifier(userName));
     }
 
     /**
@@ -1099,4 +1123,12 @@ public static String getConnectionID(
         SQLServerConnection conn = (SQLServerConnection) physicalConnection.get(pc);
         return (String) traceID.get(conn);
     }
+
+    public static void freeProcCache(Statement stmt) {
+        try {
+            stmt.execute("DBCC FREEPROCCACHE");
+        } catch (Exception e) {
+            // ignore error - some tests fails due to permission issues from managed identity, this does not seem to affect tests
+        }
+    }
 }
diff --git a/src/test/java/com/microsoft/sqlserver/jdbc/connection/PoolingTest.java b/src/test/java/com/microsoft/sqlserver/jdbc/connection/PoolingTest.java
@@ -152,6 +152,10 @@ public void testConnectionPoolClose() throws SQLException {
 
     @Test
     public void testConnectionPoolClientConnectionId() throws SQLException {
+        String auth = TestUtils.getProperty(connectionString, "authentication");
+        org.junit.Assume.assumeTrue(auth != null
+                && (auth.equalsIgnoreCase("SqlPassword") || auth.equalsIgnoreCase("ActiveDirectoryPassword")));
+
         SQLServerXADataSource ds = new SQLServerXADataSource();
         ds.setURL(connectionString);
         PooledConnection pc = null;