Move token refresh inside retry operations for Office 365 connector (#5766)

Signed-off-by: Alekhya Parisha <aparisha@amazon.com>
Co-authored-by: Alekhya Parisha <aparisha@amazon.com>

Author: alparish

data-prepper-plugins/saas-source-plugins/microsoft-office365-source/src/main/java/org/opensearch/dataprepper/plugins/source/microsoft_office365/Office365RestClient.java

@@ -75,7 +75,7 @@ public void startSubscriptions() {
         log.info("Starting Office 365 subscriptions for audit logs");
         try {
             HttpHeaders headers = new HttpHeaders();
-            headers.setBearerAuth(authConfig.getAccessToken());
+
             headers.setContentType(MediaType.APPLICATION_JSON);
 
             // TODO: Only start the subscriptions only if the call commented
@@ -103,6 +103,7 @@ public void startSubscriptions() {
 
                 RetryHandler.executeWithRetry(() -> {
                     try {
+                        headers.setBearerAuth(authConfig.getAccessToken());
                         ResponseEntity<String> response = restTemplate.exchange(
                                 url,
                                 HttpMethod.POST,
@@ -155,12 +156,13 @@ public AuditLogsResponse searchAuditLogs(final String contentType,
                         endTime.toString());
 
         final HttpHeaders headers = new HttpHeaders();
-        headers.setBearerAuth(authConfig.getAccessToken());
 
         return searchCallLatencyTimer.record(() -> {
             try {
                 return RetryHandler.executeWithRetry(
                         () -> {
+                            headers.setBearerAuth(authConfig.getAccessToken());
+
                             ResponseEntity<List<Map<String, Object>>> response = restTemplate.exchange(
                                     url,
                                     HttpMethod.GET,
@@ -205,19 +207,19 @@ public String getAuditLog(final String contentId) {
                 contentId);
 
         final HttpHeaders headers = new HttpHeaders();
-        headers.setBearerAuth(authConfig.getAccessToken());
 
         return auditLogFetchLatencyTimer.record(() -> {
             try {
-                String response = RetryHandler.executeWithRetry(() ->
-                                restTemplate.exchange(
-                                        url,
-                                        HttpMethod.GET,
-                                        new HttpEntity<>(headers),
-                                        String.class
-                                ).getBody(),
-                        authConfig::renewCredentials
-                );
+                String response = RetryHandler.executeWithRetry(() -> {
+                    headers.setBearerAuth(authConfig.getAccessToken());
+
+                    return restTemplate.exchange(
+                            url,
+                            HttpMethod.GET,
+                            new HttpEntity<>(headers),
+                            String.class
+                    ).getBody();
+                }, authConfig::renewCredentials);
                 auditLogRequestsSuccessCounter.increment();
                 return response;
             } catch (Exception e) {

data-prepper-plugins/saas-source-plugins/microsoft-office365-source/src/test/java/org/opensearch/dataprepper/plugins/source/microsoft_office365/Office365RestClientTest.java

@@ -34,6 +34,7 @@
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
+import java.util.ArrayList;
 
 import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.junit.jupiter.api.Assertions.assertDoesNotThrow;
@@ -42,6 +43,7 @@
 import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.ArgumentMatchers.anyString;
 import static org.mockito.ArgumentMatchers.eq;
+import static org.mockito.ArgumentMatchers.argThat;
 import static org.mockito.Mockito.times;
 import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;
@@ -242,4 +244,49 @@ void testGetAuditLogFailure() {
                 () -> office365RestClient.getAuditLog("test-content-id"));
         assertEquals("Failed to fetch audit log", exception.getMessage());
     }
+
+    @Test
+    void testTokenRenewal() {
+        // Setup
+        Instant startTime = Instant.now().minus(1, ChronoUnit.HOURS);
+        Instant endTime = Instant.now();
+
+        List<String> tokensUsed = new ArrayList<>();
+        List<String> requestTokens = new ArrayList<>();
+
+        when(authConfig.getTenantId()).thenReturn("test-tenant-id");
+        when(authConfig.getAccessToken()).thenAnswer(invocation -> {
+            String token = "token-" + tokensUsed.size();
+            tokensUsed.add(token);
+            return token;
+        });
+
+        when(restTemplate.exchange(
+                anyString(),
+                eq(HttpMethod.GET),
+                argThat(request -> {
+                    requestTokens.add(request.getHeaders().getFirst("Authorization"));
+                    return true;
+                }),
+                any(ParameterizedTypeReference.class)
+        )).thenAnswer(invocation -> {
+            if (requestTokens.size() == 1) {
+                throw new HttpClientErrorException(HttpStatus.UNAUTHORIZED); // First request fails with 401
+            }
+            return new ResponseEntity<>(Collections.singletonList(new HashMap<>()), HttpStatus.OK); // Second request succeeds
+        });
+
+        // Execute
+        office365RestClient.searchAuditLogs(
+                "Audit.AzureActiveDirectory",
+                startTime,
+                endTime,
+                null
+        );
+
+        // Verify
+        assertEquals(2, requestTokens.size(), "Should have made two requests");
+        assertEquals("Bearer token-0", requestTokens.get(0), "First request should use token-0");
+        assertEquals("Bearer token-1", requestTokens.get(1), "Second request should use token-1");
+    }
 }