Update fedauth tests to use cached tokens and retry if throttled

src/test/java/com/microsoft/sqlserver/jdbc/fedauth/FedauthCommon.java

@@ -159,13 +159,16 @@ public static void getConfigs() throws Exception {
      * 
      */
     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,30 +177,31 @@ static void getFedauthInfo() {
                 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());
                 }
             }
         }
     }
 
     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;

src/test/java/com/microsoft/sqlserver/jdbc/resiliency/BasicConnectionTest.java

@@ -52,8 +52,8 @@ public void testBasicReconnectDefault() throws SQLException {
     @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 void testBasicConnectionAAD() throws Exception {
                 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());
                 }
             }

src/test/java/com/microsoft/sqlserver/testframework/AbstractTest.java

@@ -32,6 +32,7 @@
 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;