diff --git a/data-prepper-plugins/saas-source-plugins/microsoft-office365-source/src/test/java/org/opensearch/dataprepper/plugins/source/microsoft_office365/Office365RestClientTest.java b/data-prepper-plugins/saas-source-plugins/microsoft-office365-source/src/test/java/org/opensearch/dataprepper/plugins/source/microsoft_office365/Office365RestClientTest.java index 2edf221303..29f51118f9 100644 --- a/data-prepper-plugins/saas-source-plugins/microsoft-office365-source/src/test/java/org/opensearch/dataprepper/plugins/source/microsoft_office365/Office365RestClientTest.java +++ b/data-prepper-plugins/saas-source-plugins/microsoft-office365-source/src/test/java/org/opensearch/dataprepper/plugins/source/microsoft_office365/Office365RestClientTest.java @@ -22,6 +22,9 @@ import org.opensearch.dataprepper.plugins.source.microsoft_office365.models.AuditLogsResponse; import org.opensearch.dataprepper.plugins.source.source_crawler.exception.SaaSCrawlerException; import org.opensearch.dataprepper.plugins.source.source_crawler.metrics.VendorAPIMetricsRecorder; +import org.opensearch.dataprepper.plugins.source.source_crawler.utils.retry.DefaultRetryStrategy; +import org.opensearch.dataprepper.plugins.source.source_crawler.utils.retry.DefaultStatusCodeHandler; +import org.opensearch.dataprepper.plugins.source.source_crawler.utils.retry.RetryHandler; import org.opensearch.dataprepper.test.helper.ReflectivelySetField; import org.springframework.core.ParameterizedTypeReference; import org.springframework.http.HttpEntity; @@ -79,6 +82,8 @@ class Office365RestClientTest { private Office365RestClient office365RestClient; + private static final int CUSTOM_MAX_RETRIES = 1; + @BeforeEach void setUp() throws NoSuchFieldException, IllegalAccessException { // Setup VendorAPIMetricsRecorder method mocks - use lenient to avoid unnecessary stubbing errors @@ -112,6 +117,12 @@ void setUp() throws NoSuchFieldException, IllegalAccessException { office365RestClient = new Office365RestClient(authConfig, pluginMetrics, metricsRecorder); ReflectivelySetField.setField(Office365RestClient.class, office365RestClient, "restTemplate", restTemplate); + + // Optionally replace RetryHandler with custom one (1 retry) for faster test execution + RetryHandler customRetryHandler = new RetryHandler( + new DefaultRetryStrategy(CUSTOM_MAX_RETRIES), + new DefaultStatusCodeHandler()); + ReflectivelySetField.setField(Office365RestClient.class, office365RestClient, "retryHandler", customRetryHandler); } /** @@ -365,11 +376,18 @@ void testGetAuditLogFailure() { * Verifies that failed authentication triggers credential renewal and retry succeeds. */ @Test - void testTokenRenewal() { + void testTokenRenewal() throws NoSuchFieldException, IllegalAccessException { // Setup Instant startTime = Instant.now().minus(1, ChronoUnit.HOURS); Instant endTime = Instant.now(); + // Override retry handler to allow at least 2 attempts for token renewal test + RetryHandler tokenRenewalRetryHandler = new RetryHandler( + new DefaultRetryStrategy(2), // Need at least 2 attempts for token renewal scenario + new DefaultStatusCodeHandler()); + ReflectivelySetField.setField(Office365RestClient.class, office365RestClient, "retryHandler", + tokenRenewalRetryHandler); + List tokensUsed = new ArrayList<>(); List requestTokens = new ArrayList<>();