fix(lambda): improve CloudWatch log polling and stabilize unit test

Author: shivamwayal37

src/main/java/io/kestra/plugin/aws/lambda/Invoke.java

@@ -210,50 +210,72 @@ Optional<String> readError(String payload) {
     }
 
     @VisibleForTesting
-    void fetchAndLogLambdaLogs(
-        RunContext runContext,
-        String functionArn,
-        Instant startTime
-    ) {
+    void fetchAndLogLambdaLogs(RunContext runContext, String functionArn, Instant startTime) {
         var logger = runContext.logger();
-
-        // Extract function name from ARN
         String functionName;
+    
         try {
-            functionName = functionArn.split(":function:")[1].split(":")[0];
+            functionName = extractFunctionName(functionArn);
         } catch (Exception e) {
             logger.warn("Unable to determine Lambda function name from ARN: {}", functionArn);
             return;
         }
 
         String logGroupName = "/aws/lambda/" + functionName;
 
+        // Polling logic: 5 attempts, waiting 3 seconds between each
+        int maxAttempts = 5;
+        long sleepMillis = 3000;
+        boolean logsFound = false;
+
         try (CloudWatchLogsClient logsClient = getCloudWatchLogsClient(runContext)) {
-            FilterLogEventsRequest request = FilterLogEventsRequest.builder()
-                .logGroupName(logGroupName)
-                .startTime(startTime.toEpochMilli())
-                .build();
-
-            // Fetch logs using CloudWatch paginator
-            logsClient.filterLogEventsPaginator(request)
-                .events()
-                .stream()
-                // Hard cap to prevent excessive log volume in a single task execution
-                .limit(1_000)
-                .map(FilteredLogEvent::message)
-                .filter(message -> message != null && !message.isBlank())
-                .forEach(message -> logger.info("[lambda] {}", message));
+            for (int i = 0; i < maxAttempts; i++) {
+                FilterLogEventsRequest request = FilterLogEventsRequest.builder()
+                    .logGroupName(logGroupName)
+                    .startTime(startTime.toEpochMilli())
+                    .build();
+
+                var response = logsClient.filterLogEvents(request);
+                var events = response.events();
 
+                if (events != null && !events.isEmpty()) {
+                    events.stream()
+                        .limit(1000)
+                        .map(FilteredLogEvent::message)
+                        .filter(message -> message != null && !message.isBlank())
+                        .forEach(message -> logger.info("[lambda] {}", message.trim()));
+                    logsFound = true;
+                    break; // Exit early if found logs
+                }
+
+                // Wait before next retry, but not after the last attempt
+                if (i < maxAttempts - 1) {
+                    Thread.sleep(sleepMillis);
+                }
+            }
+
+            if (!logsFound) {
+                logger.debug("No CloudWatch logs found for {} after {} attempts.", functionName, maxAttempts);
+            }
+
+        } catch (InterruptedException e) {
+            Thread.currentThread().interrupt();
+            logger.warn("Log fetching interrupted for Lambda {}", functionArn);
         } catch (Exception e) {
-            // Logs must never fail the task execution
-            logger.warn(
-                "Failed to fetch CloudWatch logs for Lambda {}: {}",
-                functionArn,
-                e.getMessage()
-            );
+            logger.warn("Failed to fetch CloudWatch logs for Lambda {}: {}", functionArn, e.getMessage());
         }
     }
 
+    @VisibleForTesting
+    private String extractFunctionName(String functionArnOrName) {
+        if (functionArnOrName.contains(":function:")) {
+            // Handle Full ARN
+            return functionArnOrName.split(":function:")[1].split(":")[0];
+        }
+        // Handle just the name
+        return functionArnOrName;
+    }
+
     @VisibleForTesting
     void handleError(String functionArn, ContentType contentType, SdkBytes payload) {
         String errorPayload;

src/test/java/io/kestra/plugin/aws/lambda/InvokeUnitTest.java

@@ -28,6 +28,7 @@
 import static org.mockito.Mockito.doReturn;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.spy;
+import static org.mockito.Mockito.times;
 import static org.mockito.Mockito.verify;
 import org.mockito.invocation.InvocationOnMock;
 import org.mockito.junit.jupiter.MockitoExtension;
@@ -220,37 +221,40 @@ void givenFunctionArnNoParams_whenInvokeLambda_thenOutputWithFile(
     }
 
     @Test
-    void givenLambdaInvocation_whenCloudWatchLogsEnabled_thenLogsAreFetchedAndLogged() throws Exception {
-        // Given
+    void givenLambdaInvocation_whenLogsAreFetched_thenLogsAreLoggedCorrectly() throws Exception {
+        // Arrange
+        Instant startTime = Instant.parse("2026-01-15T10:00:00Z");
+
         FilteredLogEvent logEvent = FilteredLogEvent.builder()
             .message("Hello from CloudWatch Logs")
-            .timestamp(Instant.now().toEpochMilli())
+            .timestamp(startTime.plusSeconds(1).toEpochMilli())
             .build();
 
-        // Mock paginator
-        FilterLogEventsIterable paginator = mock(FilterLogEventsIterable.class);
+        FilterLogEventsResponse response = FilterLogEventsResponse.builder()
+            .events(List.of(logEvent))
+            .build();
 
-        // Correct AWS SDK iterable for events()
-        software.amazon.awssdk.core.pagination.sync.SdkIterable<FilteredLogEvent> eventIterable =
-            () -> List.of(logEvent).iterator();
+        given(logsClient.filterLogEvents(any(FilterLogEventsRequest.class)))
+            .willReturn(response);
 
-        given(paginator.events()).willReturn(eventIterable);
-        given(logsClient.filterLogEventsPaginator(any(FilterLogEventsRequest.class)))
-            .willReturn(paginator);
+        given(context.logger()).willReturn(logger);
 
-        // Spy Invoke to inject mocked logs client
         Invoke spyInvoke = spy(invoke);
         doReturn(logsClient).when(spyInvoke).getCloudWatchLogsClient(any());
 
-        // When
+        // Act
         spyInvoke.fetchAndLogLambdaLogs(
             context,
-            "arn:aws:lambda:eu-central-1:123456789012:function:test",
-            Instant.now()
+            "arn:aws:lambda:ap-south-1:123456789012:function:test-function",
+            startTime
         );
 
-        // Then
-        verify(logger).info("[lambda] {}", "Hello from CloudWatch Logs");
-    }
+        // Assert
+        verify(logsClient, times(1))
+            .filterLogEvents(any(FilterLogEventsRequest.class));
 
+        verify(logger)
+            .info("[lambda] {}", "Hello from CloudWatch Logs");
+    }
 }
+