Repository: apache/hadoop
File: hadoop-tools/hadoop-azure/src/test/java/org/apache/hadoop/fs/azurebfs/ITestAbfsMsiTokenProvider.java (Java)

Code before:
    // Mock the getTokenSingleCall to throw exception so the retry logic comes into place.
    try (MockedStatic<AzureADAuthenticator> adAuthenticator = Mockito.mockStatic(
        AzureADAuthenticator.class, Mockito.CALLS_REAL_METHODS)) {
      adAuthenticator.when(
          () -> AzureADAuthenticator.getTokenSingleCall(Mockito.anyString(),
              Mockito.anyString(), Mockito.any(), Mockito.anyString(),
              Mockito.anyBoolean())).thenThrow(httpException);

      // Mock the tokenFetchRetryPolicy to verify retries.
      ExponentialRetryPolicy exponentialRetryPolicy = Mockito.spy(
          conf.getOauthTokenFetchRetryPolicy());
      Field tokenFetchRetryPolicy = AzureADAuthenticator.class.getDeclaredField(
          "tokenFetchRetryPolicy");
      tokenFetchRetryPolicy.setAccessible(true);
      tokenFetchRetryPolicy.set(ExponentialRetryPolicy.class,
          exponentialRetryPolicy);

      AccessTokenProvider tokenProvider = new MsiTokenProvider(authEndpoint,
          tenantGuid, clientId, authority);
      AzureADToken token = null;
      intercept(AzureADAuthenticator.HttpException.class,
          tokenProvider::getToken);

      // If the status code doesn't qualify for retry shouldRetry returns false and the loop ends.
      // It being called multiple times verifies that the retry was done for the throttling status code 429.
      Mockito.verify(exponentialRetryPolicy,
              times(DEFAULT_AZURE_OAUTH_TOKEN_FETCH_RETRY_MAX_ATTEMPTS + 1))
          .shouldRetry(Mockito.anyInt(), Mockito.anyInt());
    }
  }
}


Review comment (suggestion):
so ExponentialRetryPolicy.getRetryCount() is there to let you pass a non-mocked policy in and then assert on it. how about using that here? it probably needs making the accessors public, rather than package scoped, but that's all. The less use we make of mockito, the less things will break with every mockito upgrade

Code after:
          // Validate shouldRetry returns true until the final attempt
          if (attempt < retryPolicy.getMaxRetryCount()) {
            Assertions.assertThat(retry)
                .describedAs("Attempt %d: shouldRetry must be true for 429", attempt)
                .isTrue();
            // Simulate retry by continuing
          } else {
            // Final attempt: shouldRetry should now be false if this was last retry
            Assertions.assertThat(retry)
                .describedAs("Final attempt %d: shouldRetry can be false after max retries", attempt)
                .isTrue(); // Still true because maxRetries not exceeded yet

            // Return a valid fake token
            AzureADToken token = new AzureADToken();
            token.setAccessToken("fake-token");
            token.setExpiry(new Date(System.currentTimeMillis() + TimeUnit.HOURS.toMillis(1)));
            return token;
          }
        }
      }
    }
    AccessTokenProvider tokenProvider = new TestMsiTokenProvider(
        "https://fake-endpoint", "tenant", "clientId", "authority"
    );
    // Trigger token acquisition
    AzureADToken token = tokenProvider.getToken();
    // Assertions
    assertThat(token.getAccessToken()).isEqualTo("fake-token");
    // If the status code doesn't qualify for retry shouldRetry returns false and the loop ends.
    // It being called multiple times verifies that the retry was done for the throttling status code 429.
    Assertions.assertThat(attemptCounter.get())
        .describedAs("Number of retries should be equal to "
            + "max attempts for token fetch.")
        .isEqualTo(DEFAULT_AZURE_OAUTH_TOKEN_FETCH_RETRY_MAX_ATTEMPTS);
  }
}