Repository: microsoft/mssql-jdbc
Pull request #2077 (2023-02-17 07:44:54)

Update fedauth tests to use cached tokens and retry if throttled


Patch:
diff --git a/src/main/java/com/microsoft/sqlserver/jdbc/ISQLServerConnection.java b/src/main/java/com/microsoft/sqlserver/jdbc/ISQLServerConnection.java
index ee0a5a8e..88fa73a3 100644
--- a/src/main/java/com/microsoft/sqlserver/jdbc/ISQLServerConnection.java
+++ b/src/main/java/com/microsoft/sqlserver/jdbc/ISQLServerConnection.java
@@ -240,7 +240,7 @@ public interface ISQLServerConnection extends java.sql.Connection {
      * Valid values are: datetime, datetime2 or datetimeoffset.
      * 
      * @param datetimeParameterTypeValue
-     *        The datatype to use when encoding Java dates into SQL Server. Valid values are: 
+     *        The datatype to use when encoding Java dates into SQL Server. Valid values are:
      *        datetime, datetime2 or datetimeoffset.
      * 
      * @throws SQLServerException
@@ -451,4 +451,19 @@ public interface ISQLServerConnection extends java.sql.Connection {
      */
     @Deprecated(since = "12.1.0", forRemoval = true)
     void setMsiTokenCacheTtl(int timeToLive);
+
+    /**
+     * Returns the fully qualified class name of the implementing class for {@link SQLServerAccessTokenCallback}.
+     *
+     * @return accessTokenCallbackClass
+     */
+    String getAccessTokenCallbackClass();
+
+    /**
+     * Sets 'accessTokenCallbackClass' to the fully qualified class name
+     * of the implementing class for {@link SQLServerAccessTokenCallback}.
+     *
+     * @param accessTokenCallbackClass
+     */
+    void setAccessTokenCallbackClass(String accessTokenCallbackClass);
 }
diff --git a/src/main/java/com/microsoft/sqlserver/jdbc/ISQLServerDataSource.java b/src/main/java/com/microsoft/sqlserver/jdbc/ISQLServerDataSource.java
index 4e6c8390..a435666e 100644
--- a/src/main/java/com/microsoft/sqlserver/jdbc/ISQLServerDataSource.java
+++ b/src/main/java/com/microsoft/sqlserver/jdbc/ISQLServerDataSource.java
@@ -347,7 +347,7 @@ public interface ISQLServerDataSource extends javax.sql.CommonDataSource {
      * Sets the SQL server datatype to use for Java datetime and timestamp values.
      * 
      * @param datetimeParameterType
-     *        The SQL datatype to use when encoding Java dates for SQL Server. Valid values are: 
+     *        The SQL datatype to use when encoding Java dates for SQL Server. Valid values are:
      *        datetime, datetime2 or datetimeoffset.
      */
     void setDatetimeParameterType(String datetimeParameterType);
@@ -1268,4 +1268,19 @@ public interface ISQLServerDataSource extends javax.sql.CommonDataSource {
      * @return Access token callback delegate.
      */
     SQLServerAccessTokenCallback getAccessTokenCallback();
+
+    /**
+     * Returns the fully qualified class name of the implementing class for {@link SQLServerAccessTokenCallback}.
+     *
+     * @return accessTokenCallbackClass
+     */
+    String getAccessTokenCallbackClass();
+
+    /**
+     * Sets 'accessTokenCallbackClass' to the fully qualified class name
+     * of the implementing class for {@link SQLServerAccessTokenCallback}.
+     *
+     * @param accessTokenCallbackClass
+     */
+    void setAccessTokenCallbackClass(String accessTokenCallbackClass);
 }
diff --git a/src/main/java/com/microsoft/sqlserver/jdbc/SQLServerAccessTokenCallback.java b/src/main/java/com/microsoft/sqlserver/jdbc/SQLServerAccessTokenCallback.java
index 2d3ba451..c56c6be4 100644
--- a/src/main/java/com/microsoft/sqlserver/jdbc/SQLServerAccessTokenCallback.java
+++ b/src/main/java/com/microsoft/sqlserver/jdbc/SQLServerAccessTokenCallback.java
@@ -17,5 +17,5 @@ public interface SQLServerAccessTokenCallback {
      *
      * @return Returns a {@link SqlAuthenticationToken}.
      */
-    SqlAuthenticationToken getAccessToken(String stsurl, String spn);
+    SqlAuthenticationToken getAccessToken(String spn, String stsurl);
 }
diff --git a/src/main/java/com/microsoft/sqlserver/jdbc/SQLServerConnection.java b/src/main/java/com/microsoft/sqlserver/jdbc/SQLServerConnection.java
index db1f9513..3ee35c2a 100644
--- a/src/main/java/com/microsoft/sqlserver/jdbc/SQLServerConnection.java
+++ b/src/main/java/com/microsoft/sqlserver/jdbc/SQLServerConnection.java
@@ -144,6 +144,15 @@ public class SQLServerConnection implements ISQLServerConnection, java.io.Serial
     /** fedAuth requested flag */
     private boolean federatedAuthenticationRequested = false;
 
+    /** Access token callback class name */
+    private String accessTokenCallbackClass = null;
+
+    /** Flag that determines whether the accessToken callback class is set **/
+    private boolean hasAccessTokenCallbackClass = false;
+
+    /** Flag that determines whether the accessToken callback was set **/
+    private SQLServerAccessTokenCallback accessTokenCallback = null;
+
     /**
      * Keep this distinct from _federatedAuthenticationRequested, since some fedauth library types may not need more
      * info
@@ -539,12 +548,10 @@ public class SQLServerConnection implements ISQLServerConnection, java.io.Serial
                 default:
                     // If authenticationString not specified, check if access token callback was set.
                     // If access token callback is set, break.
-                    if (null != activeConnectionProperties
-                            .get(SQLServerDriverObjectProperty.ACCESS_TOKEN_CALLBACK.toString())) {
+                    if (null != accessTokenCallback || hasAccessTokenCallbackClass) {
                         this.authentication = SqlAuthentication.NOT_SPECIFIED;
                         break;
                     }
-                    assert (false);
                     MessageFormat form = new MessageFormat(
                             SQLServerException.getErrString("R_InvalidConnectionSetting"));
                     Object[] msgArgs = {"authentication", authenticationString};
@@ -2452,10 +2459,15 @@ public class SQLServerConnection implements ISQLServerConnection, java.io.Serial
                     ntlmAuthentication = true;
                 }
 
-                SQLServerAccessTokenCallback callback = (SQLServerAccessTokenCallback) activeConnectionProperties
+                accessTokenCallback = (SQLServerAccessTokenCallback) activeConnectionProperties
                         .get(SQLServerDriverObjectProperty.ACCESS_TOKEN_CALLBACK.toString());
 
-                if (null != callback && (!activeConnectionProperties
+                hasAccessTokenCallbackClass = null != activeConnectionProperties
+                        .get(SQLServerDriverStringProperty.ACCESS_TOKEN_CALLBACK_CLASS.toString())
+                        && !activeConnectionProperties
+                                .getProperty(SQLServerDriverStringProperty.ACCESS_TOKEN_CALLBACK_CLASS.toString())
+                                .isEmpty();
+                if ((null != accessTokenCallback || hasAccessTokenCallbackClass) && (!activeConnectionProperties
                         .getProperty(SQLServerDriverStringProperty.USER.toString()).isEmpty()
                         || !activeConnectionProperties.getProperty(SQLServerDriverStringProperty.PASSWORD.toString())
                                 .isEmpty())) {
@@ -2463,6 +2475,14 @@ public class SQLServerConnection implements ISQLServerConnection, java.io.Serial
                             SQLServerException.getErrString("R_AccessTokenCallbackWithUserPassword"), null);
                 }
 
+                sPropKey = SQLServerDriverStringProperty.ACCESS_TOKEN_CALLBACK_CLASS.toString();
+                sPropValue = activeConnectionProperties.getProperty(sPropKey);
+                if (null == sPropValue) {
+                    sPropValue = SQLServerDriverStringProperty.ACCESS_TOKEN_CALLBACK_CLASS.getDefaultValue();
+                    activeConnectionProperties.setProperty(sPropKey, sPropValue);
+                }
+                setAccessTokenCallbackClass(sPropValue);
+
                 sPropKey = SQLServerDriverStringProperty.AUTHENTICATION.toString();
                 sPropValue = activeConnectionProperties.getProperty(sPropKey);
                 if (null == sPropValue) {
@@ -3503,8 +3523,7 @@ public class SQLServerConnection implements ISQLServerConnection, java.io.Serial
     void prelogin(String serverName, int portNumber) throws SQLServerException {
         // Build a TDS Pre-Login packet to send to the server.
         if ((!authenticationString.equalsIgnoreCase(SqlAuthentication.NOT_SPECIFIED.toString()))
-                || (null != accessTokenInByte) || null != activeConnectionProperties
-                        .get(SQLServerDriverObjectProperty.ACCESS_TOKEN_CALLBACK.toString())) {
+                || (null != accessTokenInByte) || null != accessTokenCallback || hasAccessTokenCallbackClass) {
             fedAuthRequiredByUser = true;
         }
 
@@ -3886,8 +3905,8 @@ public class SQLServerConnection implements ISQLServerConnection, java.io.Serial
                     // Or AccessToken is not null, mean token based authentication is used.
                     if (((null != authenticationString)
                             && (!authenticationString.equalsIgnoreCase(SqlAuthentication.NOT_SPECIFIED.toString())))
-                            || (null != accessTokenInByte) || null != activeConnectionProperties
-                                    .get(SQLServerDriverObjectProperty.ACCESS_TOKEN_CALLBACK.toString())) {
+                            || (null != accessTokenInByte) || null != accessTokenCallback
+                            || hasAccessTokenCallbackClass) {
                         fedAuthRequiredPreLoginResponse = (preloginResponse[optionOffset] == 1);
                     }
                     break;
@@ -4830,8 +4849,7 @@ public class SQLServerConnection implements ISQLServerConnection, java.io.Serial
                             break;
                         default:
                             // If not specified, check if access token callback was set. If it is set, break.
-                            if (null != activeConnectionProperties
-                                    .get(SQLServerDriverObjectProperty.ACCESS_TOKEN_CALLBACK.toString())) {
+                            if (null != accessTokenCallback || hasAccessTokenCallbackClass) {
                                 workflow = TDS.ADALWORKFLOW_ACCESSTOKENCALLBACK;
                                 break;
                             }
@@ -5051,8 +5069,7 @@ public class SQLServerConnection implements ISQLServerConnection, java.io.Serial
                         || authenticationString
                                 .equalsIgnoreCase(SqlAuthentication.ACTIVE_DIRECTORY_INTERACTIVE.toString()))
                         && fedAuthRequiredPreLoginResponse)
-                || null != activeConnectionProperties
-                        .get(SQLServerDriverObjectProperty.ACCESS_TOKEN_CALLBACK.toString())) {
+                || null != accessTokenCallback || hasAccessTokenCallbackClass) {
             federatedAuthenticationInfoRequested = true;
             fedAuthFeatureExtensionData = new FederatedAuthenticationFeatureExtensionData(TDS.TDS_FEDAUTH_LIBRARY_ADAL,
                     authenticationString, fedAuthRequiredPreLoginResponse);
@@ -5612,11 +5629,22 @@ public class SQLServerConnection implements ISQLServerConnection, java.io.Serial
 
         attemptRefreshTokenLocked = true;
 
-        SQLServerAccessTokenCallback callback = (SQLServerAccessTokenCallback) activeConnectionProperties
-                .get(SQLServerDriverObjectProperty.ACCESS_TOKEN_CALLBACK.toString());
-
-        if (authenticationString.equals(SqlAuthentication.NOT_SPECIFIED.toString()) && null != callback) {
-            fedAuthToken = callback.getAccessToken(fedAuthInfo.spn, fedAuthInfo.stsurl);
+        if (authenticationString.equals(SqlAuthentication.NOT_SPECIFIED.toString()) && null != accessTokenCallbackClass
+                && !accessTokenCallbackClass.isEmpty()) {
+            try {
+                Object[] msgArgs = {"accessTokenCallbackClass",
+                        "com.microsoft.sqlserver.jdbc.SQLServerAccessTokenCallback"};
+                SQLServerAccessTokenCallback callbackInstance = Util.newInstance(SQLServerAccessTokenCallback.class,
+                        accessTokenCallbackClass, null, msgArgs);
+                fedAuthToken = callbackInstance.getAccessToken(fedAuthInfo.spn, fedAuthInfo.stsurl);
+            } catch (Exception e) {
+                MessageFormat form = new MessageFormat(
+                        SQLServerException.getErrString("R_InvalidAccessTokenCallbackClass"));
+                throw new SQLServerException(form.format(new Object[] {accessTokenCallbackClass}), e.getCause());
+            }
+        } else if (authenticationString.equals(SqlAuthentication.NOT_SPECIFIED.toString())
+                && null != accessTokenCallback) {
+            fedAuthToken = accessTokenCallback.getAccessToken(fedAuthInfo.spn, fedAuthInfo.stsurl);
         } else {
             fedAuthToken = getFedAuthToken(fedAuthInfo);
         }
@@ -7622,6 +7650,29 @@ public class SQLServerConnection implements ISQLServerConnection, java.io.Serial
         return 1 < getServerPreparedStatementDiscardThreshold();
     }
 
+    /**
+     * Returns the fully qualified class name of the implementing class for {@link SQLServerAccessTokenCallback}.
+     *
+     * @return accessTokenCallbackClass
+     */
+    public String getAccessTokenCallbackClass() {
+        if (null == this.accessTokenCallbackClass) {
+            return SQLServerDriverStringProperty.ACCESS_TOKEN_CALLBACK_CLASS.getDefaultValue();
+        }
+
+        return this.accessTokenCallbackClass;
+    };
+
+    /**
+     * Sets 'accessTokenCallbackClass' to the fully qualified class name
+     * of the implementing class for {@link SQLServerAccessTokenCallback}.
+     *
+     * @param accessTokenCallbackClass
+     */
+    public void setAccessTokenCallbackClass(String accessTokenCallbackClass) {
+        this.accessTokenCallbackClass = accessTokenCallbackClass;
+    };
+
     /**
      * Cleans up discarded prepared statement handles on the server using batched un-prepare actions if the batching
      * threshold has been reached.
diff --git a/src/main/java/com/microsoft/sqlserver/jdbc/SQLServerConnectionPoolProxy.java b/src/main/java/com/microsoft/sqlserver/jdbc/SQLServerConnectionPoolProxy.java
index 5e4d933c..056c7f89 100644
--- a/src/main/java/com/microsoft/sqlserver/jdbc/SQLServerConnectionPoolProxy.java
+++ b/src/main/java/com/microsoft/sqlserver/jdbc/SQLServerConnectionPoolProxy.java
@@ -661,4 +661,25 @@ class SQLServerConnectionPoolProxy implements ISQLServerConnection, java.io.Seri
     @Deprecated(since = "12.1.0", forRemoval = true)
     @Override
     public void setMsiTokenCacheTtl(int timeToLive) {}
+
+    /**
+     * Returns the fully qualified class name of the implementing class for {@link SQLServerAccessTokenCallback}.
+     *
+     * @return accessTokenCallbackClass
+     */
+    @Override
+    public String getAccessTokenCallbackClass() {
+        return wrappedConnection.getAccessTokenCallbackClass();
+    }
+
+    /**
+     * Sets 'accessTokenCallbackClass' to the fully qualified class name
+     * of the implementing class for {@link SQLServerAccessTokenCallback}.
+     *
+     * @param accessTokenCallbackClass
+     */
+    @Override
+    public void setAccessTokenCallbackClass(String accessTokenCallbackClass) {
+        wrappedConnection.setAccessTokenCallbackClass(accessTokenCallbackClass);
+    }
 }
diff --git a/src/main/java/com/microsoft/sqlserver/jdbc/SQLServerDataSource.java b/src/main/java/com/microsoft/sqlserver/jdbc/SQLServerDataSource.java
index eeb9907a..47478035 100644
--- a/src/main/java/com/microsoft/sqlserver/jdbc/SQLServerDataSource.java
+++ b/src/main/java/com/microsoft/sqlserver/jdbc/SQLServerDataSource.java
@@ -1280,6 +1280,29 @@ public class SQLServerDataSource
                 SQLServerDriverObjectProperty.ACCESS_TOKEN_CALLBACK.getDefaultValue());
     }
 
+    /**
+     * Sets 'accessTokenCallbackClass' to the fully qualified class name
+     * of the implementing class for {@link SQLServerAccessTokenCallback}.
+     *
+     * @param accessTokenCallbackClass
+     */
+    @Override
+    public void setAccessTokenCallbackClass(String accessTokenCallbackClass) {
+        setStringProperty(connectionProps, SQLServerDriverStringProperty.ACCESS_TOKEN_CALLBACK_CLASS.toString(),
+                accessTokenCallbackClass);
+    }
+
+    /**
+     * Returns the fully qualified class name of the implementing class for {@link SQLServerAccessTokenCallback}.
+     *
+     * @return accessTokenCallbackClass
+     */
+    @Override
+    public String getAccessTokenCallbackClass() {
+        return getStringProperty(connectionProps, SQLServerDriverStringProperty.ACCESS_TOKEN_CALLBACK_CLASS.toString(),
+                null);
+    }
+
     /**
      * Sets a property string value.
      *
diff --git a/src/main/java/com/microsoft/sqlserver/jdbc/SQLServerDatabaseMetaData.java b/src/main/java/com/microsoft/sqlserver/jdbc/SQLServerDatabaseMetaData.java
index f722dbe1..79f2c269 100644
--- a/src/main/java/com/microsoft/sqlserver/jdbc/SQLServerDatabaseMetaData.java
+++ b/src/main/java/com/microsoft/sqlserver/jdbc/SQLServerDatabaseMetaData.java
@@ -1667,7 +1667,7 @@ public final class SQLServerDatabaseMetaData implements java.sql.DatabaseMetaDat
                     && !name.equals(SQLServerDriverStringProperty.TRUST_STORE_PASSWORD.toString())) {
                 String val = info[index].value;
                 // skip empty strings
-                if (0 != val.length()) {
+                if (null != val && 0 != val.length()) {
                     // special case these server name, instance name and port
                     // number as these go in the front
                     if (name.equals(SQLServerDriverStringProperty.SERVER_NAME.toString())) {
diff --git a/src/main/java/com/microsoft/sqlserver/jdbc/SQLServerDriver.java b/src/main/java/com/microsoft/sqlserver/jdbc/SQLServerDriver.java
index 95b0ebca..491eac24 100644
--- a/src/main/java/com/microsoft/sqlserver/jdbc/SQLServerDriver.java
+++ b/src/main/java/com/microsoft/sqlserver/jdbc/SQLServerDriver.java
@@ -605,7 +605,8 @@ enum SQLServerDriverStringProperty {
     MAX_RESULT_BUFFER("maxResultBuffer", "-1"),
     ENCRYPT("encrypt", EncryptOption.TRUE.toString()),
     SERVER_CERTIFICATE("serverCertificate", ""),
-    DATETIME_DATATYPE("datetimeParameterType", DatetimeType.DATETIME2.toString());
+    DATETIME_DATATYPE("datetimeParameterType", DatetimeType.DATETIME2.toString()),
+    ACCESS_TOKEN_CALLBACK_CLASS("accessTokenCallbackClass", "");
 
     private final String name;
     private final String defaultValue;
@@ -833,6 +834,12 @@ public final class SQLServerDriver implements java.sql.Driver {
                     SQLServerDriverStringProperty.TRUST_MANAGER_CLASS.getDefaultValue(), false, null),
             new SQLServerDriverPropertyInfo(SQLServerDriverStringProperty.TRUST_MANAGER_CONSTRUCTOR_ARG.toString(),
                     SQLServerDriverStringProperty.TRUST_MANAGER_CONSTRUCTOR_ARG.getDefaultValue(), false, null),
+            // Callback needs to be in list despite not being settable within connection string.
+            // The reason for this is for calls to mergeURLAndSuppliedProperties to work when setting the callback.
+            new SQLServerDriverPropertyInfo(SQLServerDriverObjectProperty.ACCESS_TOKEN_CALLBACK.toString(),
+                    SQLServerDriverObjectProperty.ACCESS_TOKEN_CALLBACK.getDefaultValue(), false, null),
+            new SQLServerDriverPropertyInfo(SQLServerDriverStringProperty.ACCESS_TOKEN_CALLBACK_CLASS.toString(),
+                    SQLServerDriverStringProperty.ACCESS_TOKEN_CALLBACK_CLASS.getDefaultValue(), false, null),
             new SQLServerDriverPropertyInfo(SQLServerDriverBooleanProperty.REPLICATION.toString(),
                     Boolean.toString(SQLServerDriverBooleanProperty.REPLICATION.getDefaultValue()), false, TRUE_FALSE),
             new SQLServerDriverPropertyInfo(SQLServerDriverBooleanProperty.SEND_TIME_AS_DATETIME.toString(),
@@ -931,6 +938,8 @@ public final class SQLServerDriver implements java.sql.Driver {
     private static final SQLServerDriverPropertyInfo[] DRIVER_PROPERTIES_PROPERTY_ONLY = {
             // default required available choices
             // property name value property (if appropriate)
+            new SQLServerDriverPropertyInfo(SQLServerDriverObjectProperty.ACCESS_TOKEN_CALLBACK.toString(),
+                    SQLServerDriverObjectProperty.ACCESS_TOKEN_CALLBACK.getDefaultValue(), false, null),
             new SQLServerDriverPropertyInfo(SQLServerDriverStringProperty.ACCESS_TOKEN.toString(),
                     SQLServerDriverStringProperty.ACCESS_TOKEN.getDefaultValue(), false, null),
             new SQLServerDriverPropertyInfo(SQLServerDriverObjectProperty.GSS_CREDENTIAL.toString(),
@@ -1074,6 +1083,9 @@ public final class SQLServerDriver implements java.sql.Driver {
                     fixedup.setProperty(newname, val);
                 } else if ("gsscredential".equalsIgnoreCase(newname) && (props.get(name) instanceof GSSCredential)) {
                     fixedup.put(newname, props.get(name));
+                } else if ("accessTokenCallback".equalsIgnoreCase(newname)
+                        && (props.get(name) instanceof SQLServerAccessTokenCallback)) {
+                    fixedup.put(newname, props.get(name));
                 } else {
                     MessageFormat form = new MessageFormat(SQLServerException.getErrString("R_invalidpropertyValue"));
                     Object[] msgArgs = {name};
diff --git a/src/main/java/com/microsoft/sqlserver/jdbc/SQLServerResource.java b/src/main/java/com/microsoft/sqlserver/jdbc/SQLServerResource.java
index 48900ada..c936a572 100644
--- a/src/main/java/com/microsoft/sqlserver/jdbc/SQLServerResource.java
+++ b/src/main/java/com/microsoft/sqlserver/jdbc/SQLServerResource.java
@@ -240,6 +240,8 @@ public final class SQLServerResource extends ListResourceBundle {
         {"R_delayLoadingLobsPropertyDescription", "Boolean flag which indicates whether the driver will load LOB datatypes into memory."},
         {"R_AADSecurePrincipalIdPropertyDescription", "The Application Id of a registered application which has been granted permission to the database connected."},
         {"R_AADSecurePrincipalSecretPropertyDescription", "A Secret defined for a registered application which has been granted permission to the database connected."},
+        {"R_accessTokenCallbackClassPropertyDescription", "The class to instantiate as the SQLServerAccessTokenCallback for acquiring tokens."},
+        {"R_accessTokenCallbackPropertyDescription", "A SQLServerAccessTokenCallback object which is used to call a callback method to return an access token."},
         {"R_noParserSupport", "An error occurred while instantiating the required parser. Error: \"{0}\""},
         {"R_writeOnlyXML", "Cannot read from this SQLXML instance. This instance is for writing data only."},
         {"R_dataHasBeenReadXML", "Cannot read from this SQLXML instance. The data has already been read."},
@@ -414,6 +416,7 @@ public final class SQLServerResource extends ListResourceBundle {
         {"R_AECertHashEmpty", "Internal error. Empty certificate thumbprint specified in certificate path {0}."},
         {"R_AECertNotFound", "Certificate with thumbprint {2} not found in certificate store {1} in certificate location {0}. Verify the certificate path in the column master key definition in the database is correct, and the certificate has been imported correctly into the certificate location/store."},
         {"R_AEMaloc", "Memory allocation failure."},
+        {"R_InvalidAccessTokenCallbackClass", "Invalid accessTokenCallbackClass: {0}"},
         {"R_AEKeypathLong", "Internal error. Specified certificate path has {0} bytes, which exceeds maximum length of {1} bytes."},
         {"R_AEECEKLenBad", "The specified encrypted column encryption key''s ciphertext length: {0} does not match the ciphertext length: {1} when using column master key (certificate) in \"{2}\". The encrypted column encryption key may be corrupt, or the specified certificate path may be incorrect."},
         {"R_AEECEKSigLenBad", "The specified encrypted column encryption key''s signature length {0} does not match the length {1} when using the column master key (certificate) in \"{2}\". The encrypted column encryption key may be corrupt, or the specified certificate path may be incorrect."},


Test patch:
diff --git a/src/test/java/com/microsoft/sqlserver/jdbc/SQLServerConnectionTest.java b/src/test/java/com/microsoft/sqlserver/jdbc/SQLServerConnectionTest.java
index 5b4b2a51..f1e2298c 100644
--- a/src/test/java/com/microsoft/sqlserver/jdbc/SQLServerConnectionTest.java
+++ b/src/test/java/com/microsoft/sqlserver/jdbc/SQLServerConnectionTest.java
@@ -4,7 +4,6 @@
  */
 package com.microsoft.sqlserver.jdbc;
 
-import static org.junit.Assert.assertNotEquals;
 import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.junit.jupiter.api.Assertions.assertTrue;
 import static org.junit.jupiter.api.Assertions.fail;
@@ -1108,68 +1107,4 @@ public class SQLServerConnectionTest extends AbstractTest {
             assertTrue(e.getMessage().matches(TestUtils.formatErrorMsg("R_errorServerName")));
         }
     }
-
-    @Tag(Constants.xSQLv11)
-    @Tag(Constants.xSQLv12)
-    @Tag(Constants.xSQLv14)
-    @Tag(Constants.xSQLv15)
-    @Tag(Constants.xAzureSQLDW)
-    @Tag(Constants.reqExternalSetup)
-    @Test
-    public void testDSPooledConnectionAccessTokenCallback() throws Exception {
-        SQLServerConnectionPoolDataSource ds = new SQLServerConnectionPoolDataSource();
-
-        // User/password is not required for access token callback
-        String cs = TestUtils.addOrOverrideProperty(connectionString, "user", "");
-        cs = TestUtils.addOrOverrideProperty(cs, "password", "");
-        AbstractTest.updateDataSource(cs, ds);
-        ds.setAccessTokenCallback(TestUtils.accessTokenCallback);
-
-        TestUtils.expireTokenToggle = false;
-        SQLServerPooledConnection pc = (SQLServerPooledConnection) ds.getPooledConnection();
-        String conn1ID;
-        String conn2ID;
-
-        // Callback should provide valid token on connection open for all new connections
-        // When the access token hasn't expired, the connection ID should be the same
-        try (Connection conn1 = pc.getConnection()) {}
-        conn1ID = TestUtils.getConnectionID(pc);
-        try (Connection conn2 = pc.getConnection()) {}
-        conn2ID = TestUtils.getConnectionID(pc);
-        assertEquals(conn1ID, conn2ID);
-    }
-
-    @Tag(Constants.xSQLv11)
-    @Tag(Constants.xSQLv12)
-    @Tag(Constants.xSQLv14)
-    @Tag(Constants.xSQLv15)
-    @Tag(Constants.xAzureSQLDW)
-    @Tag(Constants.reqExternalSetup)
-    @Test
-    public void testDSPooledConnectionAccessTokenCallbackExpiredToken() throws Exception {
-        SQLServerConnectionPoolDataSource ds = new SQLServerConnectionPoolDataSource();
-
-        // User/password is not required for access token callback
-        String cs = TestUtils.addOrOverrideProperty(connectionString, "user", "");
-        cs = TestUtils.addOrOverrideProperty(cs, "password", "");
-        AbstractTest.updateDataSource(cs, ds);
-        ds.setAccessTokenCallback(TestUtils.accessTokenCallback);
-
-        SQLServerPooledConnection pc = (SQLServerPooledConnection) ds.getPooledConnection();
-        String conn1ID;
-        String conn2ID;
-
-        // When token expires after first connection, it should create a new connection to get a new token.
-        // Connection ID should not be the same.
-        TestUtils.expireTokenToggle = true;
-        pc = (SQLServerPooledConnection) ds.getPooledConnection();
-        try (Connection conn1 = pc.getConnection()) {}
-        conn1ID = TestUtils.getConnectionID(pc);
-        // Sleep until token expires
-        Thread.sleep(720000);
-        try (Connection conn2 = pc.getConnection()) {}
-        conn2ID = TestUtils.getConnectionID(pc);
-
-        assertNotEquals(conn1ID, conn2ID);
-    }
 }
diff --git a/src/test/java/com/microsoft/sqlserver/jdbc/TestUtils.java b/src/test/java/com/microsoft/sqlserver/jdbc/TestUtils.java
index 88ad138f..2f20821a 100644
--- a/src/test/java/com/microsoft/sqlserver/jdbc/TestUtils.java
+++ b/src/test/java/com/microsoft/sqlserver/jdbc/TestUtils.java
@@ -96,6 +96,8 @@ public final class TestUtils {
 
     public static final int TEST_TOKEN_EXPIRY_SECONDS = 120; // token expiry time in secs
 
+    public static String ACCESS_TOKEN_CALLBACK = null;
+
     static String applicationKey;
     static String applicationClientID;
 
@@ -132,9 +134,11 @@ public final class TestUtils {
                 String accessToken = authenticationResult.accessToken();
                 long expiresOn = authenticationResult.expiresOnDate().getTime();
 
+                ACCESS_TOKEN_CALLBACK = accessToken;
+
                 if (expireTokenToggle) {
                     Date now = new Date();
-                    long minutesToExpireWithin = 10 * 60 * 1000; // Expire within 10 minutes
+                    long minutesToExpireWithin = TEST_TOKEN_EXPIRY_SECONDS * 1000; // Expire within 2 minutes
                     return new SqlAuthenticationToken(accessToken, now.getTime() + minutesToExpireWithin);
                 } else {
                     return new SqlAuthenticationToken(accessToken, expiresOn);
@@ -161,6 +165,26 @@ public final class TestUtils {
         }
     }
 
+    public static void setAccessTokenExpiry(Object con) {
+        Field fedAuthTokenField;
+        Field wrappedConnection;
+        try {
+            fedAuthTokenField = SQLServerConnection.class.getDeclaredField("fedAuthToken");
+            fedAuthTokenField.setAccessible(true);
+
+            wrappedConnection = SQLServerConnectionPoolProxy.class.getDeclaredField("wrappedConnection");
+            wrappedConnection.setAccessible(true);
+            Object wrappedConnectionObj = wrappedConnection.get(con);
+
+            Date newExpiry = new Date(
+                    System.currentTimeMillis() + TimeUnit.SECONDS.toMillis(TEST_TOKEN_EXPIRY_SECONDS));
+            SqlAuthenticationToken newFedAuthToken = new SqlAuthenticationToken(ACCESS_TOKEN_CALLBACK, newExpiry);
+            fedAuthTokenField.set(wrappedConnectionObj, newFedAuthToken);
+        } catch (NoSuchFieldException | SecurityException | IllegalArgumentException | IllegalAccessException e) {
+            Assert.fail("Failed to set token expiry: " + e.getMessage());
+        }
+    }
+
     private TestUtils() {}
 
     /**
diff --git a/src/test/java/com/microsoft/sqlserver/jdbc/connection/RequestBoundaryMethodsTest.java b/src/test/java/com/microsoft/sqlserver/jdbc/connection/RequestBoundaryMethodsTest.java
index daba4459..d0ce2556 100644
--- a/src/test/java/com/microsoft/sqlserver/jdbc/connection/RequestBoundaryMethodsTest.java
+++ b/src/test/java/com/microsoft/sqlserver/jdbc/connection/RequestBoundaryMethodsTest.java
@@ -502,6 +502,8 @@ public class RequestBoundaryMethodsTest extends AbstractTest {
         verifiedMethodNames.add("setIPAddressPreference");
         verifiedMethodNames.add("getMsiTokenCacheTtl");
         verifiedMethodNames.add("setMsiTokenCacheTtl");
+        verifiedMethodNames.add("getAccessTokenCallbackClass");
+        verifiedMethodNames.add("setAccessTokenCallbackClass");
         return verifiedMethodNames;
     }
 }
diff --git a/src/test/java/com/microsoft/sqlserver/jdbc/fedauth/FedauthCommon.java b/src/test/java/com/microsoft/sqlserver/jdbc/fedauth/FedauthCommon.java
index 41beff62..a76bcaf5 100644
--- a/src/test/java/com/microsoft/sqlserver/jdbc/fedauth/FedauthCommon.java
+++ b/src/test/java/com/microsoft/sqlserver/jdbc/fedauth/FedauthCommon.java
@@ -159,13 +159,16 @@ public class FedauthCommon extends AbstractTest {
      * 
      */
     static void getFedauthInfo() {
-        int retry = THROTTLE_RETRY_COUNT;
+        int retry = 0;
         long interval = THROTTLE_RETRY_INTERVAL;
-        while (retry > 0) {
+        while (retry <= THROTTLE_RETRY_COUNT) {
             try {
-                final PublicClientApplication clientApplication = PublicClientApplication.builder(fedauthClientId)
-                        .executorService(Executors.newFixedThreadPool(1)).authority(stsurl).build();
-                final CompletableFuture<IAuthenticationResult> future = clientApplication
+                if (null == fedauthPcaApp) {
+                    fedauthPcaApp = PublicClientApplication.builder(fedauthClientId)
+                            .executorService(Executors.newFixedThreadPool(1)).authority(stsurl).build();
+                }
+
+                final CompletableFuture<IAuthenticationResult> future = fedauthPcaApp
                         .acquireToken(UserNamePasswordParameters.builder(Collections.singleton(spn + "/.default"),
                                 azureUserName, azurePassword.toCharArray()).build());
 
@@ -174,14 +177,14 @@ public class FedauthCommon extends AbstractTest {
                 secondsBeforeExpiration = TimeUnit.MILLISECONDS
                         .toSeconds(authenticationResult.expiresOnDate().getTime() - new Date().getTime());
                 accessToken = authenticationResult.accessToken();
-                retry = 0;
+                retry = THROTTLE_RETRY_COUNT + 1;
             } catch (MsalThrottlingException te) {
-                interval = ((MsalThrottlingException) te).retryInMs();
-                if (!checkForRetry(te, retry--, interval)) {
+                interval = te.retryInMs();
+                if (!checkForRetry(te, retry++, interval)) {
                     fail(ERR_FAILED_FEDAUTH + "no more retries: " + te.getMessage());
                 }
             } catch (Exception e) {
-                if (!checkForRetry(e, retry--, interval)) {
+                if (!checkForRetry(e, retry++, interval)) {
                     fail(ERR_FAILED_FEDAUTH + "no more retries: " + e.getMessage());
                 }
             }
@@ -189,15 +192,16 @@ public class FedauthCommon extends AbstractTest {
     }
 
     static boolean checkForRetry(Exception e, int retry, long interval) {
-        if (retry <= 0) {
+        if (retry > THROTTLE_RETRY_COUNT) {
             return false;
         }
         try {
-            System.out.println(e.getMessage() + "Get FedAuth token failed retry #" + retry + " in " + interval + " ms");
-            e.printStackTrace();
-
+            System.out
+                    .println(e.getMessage() + "Get FedAuth token failed, retry #" + retry + " in " + interval + " ms");
             Thread.sleep(interval);
         } catch (InterruptedException ex) {
+            e.printStackTrace();
+
             fail(ERR_FAILED_FEDAUTH + ex.getMessage());
         }
         return true;
diff --git a/src/test/java/com/microsoft/sqlserver/jdbc/fedauth/FedauthTest.java b/src/test/java/com/microsoft/sqlserver/jdbc/fedauth/FedauthTest.java
index 53c3d546..fc398312 100644
--- a/src/test/java/com/microsoft/sqlserver/jdbc/fedauth/FedauthTest.java
+++ b/src/test/java/com/microsoft/sqlserver/jdbc/fedauth/FedauthTest.java
@@ -40,9 +40,12 @@ public class FedauthTest extends FedauthCommon {
     static String charTable = TestUtils
             .escapeSingleQuotes(AbstractSQLGenerator.escapeIdentifier(RandomUtil.getIdentifier("JDBC_FedAuthTest")));
 
+    private static String accessTokenCallbackConnectionString;
+
     @BeforeAll
     public static void setupTests() throws Exception {
         connectionString = TestUtils.addOrOverrideProperty(connectionString, "trustServerCertificate", "true");
+        accessTokenCallbackConnectionString = "jdbc:sqlserver://" + azureServer + ";database=" + azureDatabase + ";";
         setConnection();
     }
 
@@ -352,6 +355,18 @@ public class FedauthTest extends FedauthCommon {
         validateException(url, "R_BothUserPasswordandDeprecated");
     }
 
+    @Tag(Constants.xSQLv11)
+    @Tag(Constants.xSQLv12)
+    @Tag(Constants.xSQLv14)
+    @Tag(Constants.xAzureSQLDW)
+    @Tag(Constants.reqExternalSetup)
+    @Test
+    public void testAccessTokenCallbackClassConnection() throws Exception {
+        String cs = TestUtils.addOrOverrideProperty(accessTokenCallbackConnectionString, "accessTokenCallbackClass",
+                PooledConnectionTest.AccessTokenCallbackClass.class.getName());
+        try (Connection conn1 = DriverManager.getConnection(cs)) {}
+    }
+
     private static void validateException(String url, String resourceKey) {
         try (Connection conn = DriverManager.getConnection(url)) {
             fail(TestResource.getResource("R_expectedFailPassed"));
diff --git a/src/test/java/com/microsoft/sqlserver/jdbc/fedauth/PooledConnectionTest.java b/src/test/java/com/microsoft/sqlserver/jdbc/fedauth/PooledConnectionTest.java
index 55415bbe..e109f0d6 100644
--- a/src/test/java/com/microsoft/sqlserver/jdbc/fedauth/PooledConnectionTest.java
+++ b/src/test/java/com/microsoft/sqlserver/jdbc/fedauth/PooledConnectionTest.java
@@ -4,6 +4,8 @@
  */
 package com.microsoft.sqlserver.jdbc.fedauth;
 
+import static org.junit.Assert.assertNotEquals;
+import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertTrue;
 import static org.junit.jupiter.api.Assertions.fail;
 
@@ -13,9 +15,13 @@ import java.sql.DriverManager;
 import java.sql.SQLException;
 import java.sql.Statement;
 import java.util.ArrayList;
+import java.util.HashSet;
 import java.util.List;
 import java.util.Random;
+import java.util.Set;
+import java.util.UUID;
 import java.util.concurrent.Callable;
+import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.Executors;
 import java.util.concurrent.Future;
@@ -23,6 +29,16 @@ import java.util.concurrent.TimeUnit;
 
 import javax.sql.PooledConnection;
 
+import com.microsoft.aad.msal4j.ClientCredentialFactory;
+import com.microsoft.aad.msal4j.ClientCredentialParameters;
+import com.microsoft.aad.msal4j.ConfidentialClientApplication;
+import com.microsoft.aad.msal4j.IAuthenticationResult;
+import com.microsoft.aad.msal4j.IClientCredential;
+import com.microsoft.sqlserver.jdbc.SQLServerAccessTokenCallback;
+import com.microsoft.sqlserver.jdbc.SQLServerException;
+import com.microsoft.sqlserver.jdbc.SqlAuthenticationToken;
+import com.microsoft.sqlserver.jdbc.TestResource;
+import com.microsoft.sqlserver.testframework.AbstractTest;
 import org.junit.jupiter.api.AfterAll;
 import org.junit.jupiter.api.BeforeAll;
 import org.junit.jupiter.api.Tag;
@@ -46,9 +62,40 @@ public class PooledConnectionTest extends FedauthCommon {
     static String charTable = TestUtils.escapeSingleQuotes(
             AbstractSQLGenerator.escapeIdentifier(RandomUtil.getIdentifier("JDBC_PooledConnection")));
 
+    private static String accessTokenCallbackConnectionString;
+
+    public static class AccessTokenCallbackClass implements SQLServerAccessTokenCallback {
+        @Override
+        public SqlAuthenticationToken getAccessToken(String spn, String stsurl) {
+            String scope = spn + "/.default";
+            Set<String> scopes = new HashSet<>();
+            scopes.add(scope);
+
+            try {
+                ExecutorService executorService = Executors.newSingleThreadExecutor();
+                IClientCredential credential = ClientCredentialFactory.createFromSecret(applicationKey);
+                ConfidentialClientApplication clientApplication = ConfidentialClientApplication
+                        .builder(applicationClientID, credential).executorService(executorService).authority(stsurl)
+                        .build();
+                CompletableFuture<IAuthenticationResult> future = clientApplication
+                        .acquireToken(ClientCredentialParameters.builder(scopes).build());
+
+                IAuthenticationResult authenticationResult = future.get();
+                String accessToken = authenticationResult.accessToken();
+                long expiresOn = authenticationResult.expiresOnDate().getTime();
+
+                return new SqlAuthenticationToken(accessToken, expiresOn);
+            } catch (Exception e) {
+                fail(TestResource.getResource("R_unexpectedException") + e.getMessage());
+            }
+            return null;
+        }
+    }
+
     @BeforeAll
     public static void setupTests() throws Exception {
         setConnection();
+        accessTokenCallbackConnectionString = "jdbc:sqlserver://" + azureServer + ";database=" + azureDatabase + ";";
     }
 
     @Test
@@ -263,6 +310,167 @@ public class PooledConnectionTest extends FedauthCommon {
         }
     }
 
+    @Tag(Constants.xSQLv11)
+    @Tag(Constants.xSQLv12)
+    @Tag(Constants.xSQLv14)
+    @Tag(Constants.xAzureSQLDW)
+    @Tag(Constants.reqExternalSetup)
+    @Test
+    public void testDSPooledConnectionAccessTokenCallback() throws Exception {
+        SQLServerConnectionPoolDataSource ds = new SQLServerConnectionPoolDataSource();
+
+        // User/password is not required for access token callback
+        AbstractTest.updateDataSource(accessTokenCallbackConnectionString, ds);
+        ds.setAccessTokenCallback(TestUtils.accessTokenCallback);
+
+        TestUtils.expireTokenToggle = false;
+        SQLServerPooledConnection pc = (SQLServerPooledConnection) ds.getPooledConnection();
+        String conn1ID;
+        String conn2ID;
+
+        // Callback should provide valid token on connection open for all new connections
+        // When the access token hasn't expired, the connection ID should be the same
+        try (Connection conn1 = pc.getConnection()) {}
+        conn1ID = TestUtils.getConnectionID(pc);
+        try (Connection conn2 = pc.getConnection()) {}
+        conn2ID = TestUtils.getConnectionID(pc);
+        assertEquals(conn1ID, conn2ID);
+    }
+
+    @Tag(Constants.xSQLv11)
+    @Tag(Constants.xSQLv12)
+    @Tag(Constants.xSQLv14)
+    @Tag(Constants.xAzureSQLDW)
+    @Tag(Constants.reqExternalSetup)
+    @Test
+    public void testDSPooledConnectionMergeAccessTokenCallbackProperty() throws Exception {
+        SQLServerConnectionPoolDataSource ds = new SQLServerConnectionPoolDataSource();
+
+        // Access token callback property from ds.setAccessTokenCallback should merge supplied properties from URL
+        ds.setURL(accessTokenCallbackConnectionString);
+        ds.setAccessTokenCallback(TestUtils.accessTokenCallback);
+
+        SQLServerPooledConnection pc = (SQLServerPooledConnection) ds.getPooledConnection();
+        String conn1ID;
+        String conn2ID;
+
+        try (Connection conn1 = pc.getConnection()) {}
+        conn1ID = TestUtils.getConnectionID(pc);
+        try (Connection conn2 = pc.getConnection()) {}
+        conn2ID = TestUtils.getConnectionID(pc);
+        assertEquals(conn1ID, conn2ID);
+    }
+
+    @Tag(Constants.xSQLv11)
+    @Tag(Constants.xSQLv12)
+    @Tag(Constants.xSQLv14)
+    @Tag(Constants.xAzureSQLDW)
+    @Tag(Constants.reqExternalSetup)
+    @Test
+    public void testDSPooledConnectionAccessTokenCallbackExpiredToken() throws Exception {
+        SQLServerConnectionPoolDataSource ds = new SQLServerConnectionPoolDataSource();
+
+        // User/password is not required for access token callback
+        AbstractTest.updateDataSource(accessTokenCallbackConnectionString, ds);
+        ds.setAccessTokenCallback(TestUtils.accessTokenCallback);
+
+        SQLServerPooledConnection pc = (SQLServerPooledConnection) ds.getPooledConnection();
+        String conn1ID;
+        String conn2ID;
+
+        // When token expires after first connection, it should create a new connection to get a new token.
+        // Connection ID should not be the same.
+        TestUtils.expireTokenToggle = true;
+        pc = (SQLServerPooledConnection) ds.getPooledConnection();
+        try (Connection conn1 = pc.getConnection()) {
+            TestUtils.setAccessTokenExpiry(conn1);
+        }
+        conn1ID = TestUtils.getConnectionID(pc);
+        // Sleep until token expires
+        Thread.sleep(130 * 1000);
+        try (Connection conn2 = pc.getConnection()) {}
+        conn2ID = TestUtils.getConnectionID(pc);
+
+        assertNotEquals(conn1ID, conn2ID);
+    }
+
+    @Tag(Constants.xSQLv11)
+    @Tag(Constants.xSQLv12)
+    @Tag(Constants.xSQLv14)
+    @Tag(Constants.xAzureSQLDW)
+    @Tag(Constants.reqExternalSetup)
+    @Test
+    public void testDSPooledConnectionAccessTokenCallbackClass() throws Exception {
+        SQLServerConnectionPoolDataSource ds = new SQLServerConnectionPoolDataSource();
+
+        // User/password is not required for access token callback
+        AbstractTest.updateDataSource(accessTokenCallbackConnectionString, ds);
+        ds.setAccessTokenCallbackClass(AccessTokenCallbackClass.class.getName());
+
+        SQLServerPooledConnection pc = (SQLServerPooledConnection) ds.getPooledConnection();
+        String conn1ID;
+        String conn2ID;
+
+        // Callback should provide valid token on connection open for all new connections
+        // When the access token hasn't expired, the connection ID should be the same
+        try (Connection conn1 = pc.getConnection()) {}
+        conn1ID = TestUtils.getConnectionID(pc);
+        try (Connection conn2 = pc.getConnection()) {}
+        conn2ID = TestUtils.getConnectionID(pc);
+        assertEquals(conn1ID, conn2ID);
+    }
+
+    @Tag(Constants.xSQLv11)
+    @Tag(Constants.xSQLv12)
+    @Tag(Constants.xSQLv14)
+    @Tag(Constants.xAzureSQLDW)
+    @Tag(Constants.reqExternalSetup)
+    @Test
+    public void testDSPooledConnectionAccessTokenCallbackClassExceptions() throws Exception {
+        SQLServerConnectionPoolDataSource ds = new SQLServerConnectionPoolDataSource();
+
+        // User/password is not required for access token callback
+        AbstractTest.updateDataSource(accessTokenCallbackConnectionString, ds);
+        ds.setAccessTokenCallbackClass(AccessTokenCallbackClass.class.getName());
+        ds.setUser("user");
+        SQLServerPooledConnection pc;
+
+        // Should fail with user set
+        try {
+            pc = (SQLServerPooledConnection) ds.getPooledConnection();
+            fail(TestResource.getResource("R_expectedFailPassed"));
+        } catch (SQLServerException e) {
+            assertTrue(e.getMessage().matches(TestUtils.formatErrorMsg("R_AccessTokenCallbackWithUserPassword")));
+        }
+
+        ds.setUser("");
+        ds.setPassword(UUID.randomUUID().toString());
+
+        // Should fail with password set
+        try {
+            pc = (SQLServerPooledConnection) ds.getPooledConnection();
+            fail(TestResource.getResource("R_expectedFailPassed"));
+        } catch (SQLServerException e) {
+            assertTrue(e.getMessage().matches(TestUtils.formatErrorMsg("R_AccessTokenCallbackWithUserPassword")));
+        }
+
+        // Should fail with invalid accessTokenCallbackClass value
+        ds.setAccessTokenCallbackClass("Invalid");
+        ds.setUser("");
+        ds.setPassword("");
+        try {
+            pc = (SQLServerPooledConnection) ds.getPooledConnection();
+            fail(TestResource.getResource("R_expectedFailPassed"));
+        } catch (SQLServerException e) {
+            assertTrue(e.getMessage().matches(TestUtils.formatErrorMsg("R_InvalidAccessTokenCallbackClass")));
+        }
+
+        // Should pass with no user or password set
+        ds.setAccessTokenCallbackClass(AccessTokenCallbackClass.class.getName());
+        pc = (SQLServerPooledConnection) ds.getPooledConnection();
+        try (Connection conn1 = pc.getConnection()) {}
+    }
+
     @AfterAll
     public static void terminate() throws SQLException {
         try (Connection conn = DriverManager.getConnection(adPasswordConnectionStr);
diff --git a/src/test/java/com/microsoft/sqlserver/jdbc/resiliency/BasicConnectionTest.java b/src/test/java/com/microsoft/sqlserver/jdbc/resiliency/BasicConnectionTest.java
index 52f08dfd..6b7d5c5d 100644
--- a/src/test/java/com/microsoft/sqlserver/jdbc/resiliency/BasicConnectionTest.java
+++ b/src/test/java/com/microsoft/sqlserver/jdbc/resiliency/BasicConnectionTest.java
@@ -52,8 +52,8 @@ public class BasicConnectionTest extends AbstractTest {
     @Tag(Constants.fedAuth)
     public void testBasicConnectionAAD() throws Exception {
         // retry since this could fail due to server throttling
-        int retry = THROTTLE_RETRY_COUNT;
-        while (retry > 0) {
+        int retry = 1;
+        while (retry <= THROTTLE_RETRY_COUNT) {
             try {
                 String azureServer = getConfiguredProperty("azureServer");
                 String azureDatabase = getConfiguredProperty("azureDatabase");
@@ -64,15 +64,17 @@ public class BasicConnectionTest extends AbstractTest {
                 basicReconnect("jdbc:sqlserver://" + azureServer + ";database=" + azureDatabase + ";user="
                         + azureUserName + ";password=" + azurePassword
                         + ";loginTimeout=90;Authentication=ActiveDirectoryPassword");
+		retry = THROTTLE_RETRY_COUNT + 1;
             } catch (Exception e) {
                 if (e.getMessage().matches(TestUtils.formatErrorMsg("R_crClientAllRecoveryAttemptsFailed"))) {
-                    System.out.println(e.getMessage() + "Recovery failed retry #" + retry + " in "
+                    System.out.println(e.getMessage() + ". Recovery failed, retry #" + retry + " in "
                             + THROTTLE_RETRY_INTERVAL + " ms");
-                    e.printStackTrace();
 
                     Thread.sleep(THROTTLE_RETRY_INTERVAL);
-                    retry--;
+                    retry++;
                 } else {
+                    e.printStackTrace();
+
                     fail(e.getMessage());
                 }
             }
diff --git a/src/test/java/com/microsoft/sqlserver/testframework/AbstractTest.java b/src/test/java/com/microsoft/sqlserver/testframework/AbstractTest.java
index a504f9ca..c3656186 100644
--- a/src/test/java/com/microsoft/sqlserver/testframework/AbstractTest.java
+++ b/src/test/java/com/microsoft/sqlserver/testframework/AbstractTest.java
@@ -32,6 +32,7 @@ import org.junit.jupiter.api.AfterAll;
 import org.junit.jupiter.api.Assertions;
 import org.junit.jupiter.api.BeforeAll;
 
+import com.microsoft.aad.msal4j.PublicClientApplication;
 import com.microsoft.sqlserver.jdbc.ISQLServerDataSource;
 import com.microsoft.sqlserver.jdbc.SQLServerColumnEncryptionAzureKeyVaultProvider;
 import com.microsoft.sqlserver.jdbc.SQLServerColumnEncryptionJavaKeyStoreProvider;
@@ -101,6 +102,8 @@ public abstract class AbstractTest {
     protected static String connectionString = null;
     protected static String connectionStringNTLM;
 
+    protected static PublicClientApplication fedauthPcaApp = null;
+
     private static boolean determinedSqlAzureOrSqlServer = false;
     private static boolean determinedSqlOS = false;
     private static boolean isSqlAzure = false;