Optimized RestClient Tests using customized Retry Handler

Signed-off-by: eatulban <[email protected]>

Author: eatulban

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.auth.Office365AuthenticationInterface;
 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.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;
@@ -64,10 +67,18 @@ class Office365RestClientTest {
 
     private Office365RestClient office365RestClient;
 
+    private static final int CUSTOM_MAX_RETRIES = 1;
+
     @BeforeEach
     void setUp() throws NoSuchFieldException, IllegalAccessException{
         office365RestClient = new Office365RestClient(authConfig, pluginMetrics);
         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);
     }
 
     @Test
@@ -255,11 +266,18 @@ void testGetAuditLogFailure() {
     }
 
     @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<String> tokensUsed = new ArrayList<>();
         List<String> requestTokens = new ArrayList<>();
 
@@ -523,8 +541,8 @@ void testApiCallsCounterIncrementOnRetries() throws NoSuchFieldException, Illega
         // Execute and expect exception
         assertThrows(RuntimeException.class, () -> office365RestClient.getAuditLog(contentUri));
 
-        // Verify apiCallsCounter was incremented 6 times (once for each retry attempt)
-        verify(mockApiCallsCounter, times(6)).increment();
+        // Verify apiCallsCounter was incremented CUSTOM_MAX_RETRIES times
+        verify(mockApiCallsCounter, times(CUSTOM_MAX_RETRIES)).increment();
     }
 
     @Test
@@ -549,7 +567,7 @@ void testApiCallsCounterIncrementForSearchAuditLogsWithRetries() throws NoSuchFi
                 null
         ));
 
-        // Verify apiCallsCounter was incremented 6 times (once for each retry attempt)
-        verify(mockApiCallsCounter, times(6)).increment();
+        // Verify apiCallsCounter was incremented CUSTOM_MAX_RETRIES times
+        verify(mockApiCallsCounter, times(CUSTOM_MAX_RETRIES)).increment();
     }
 }