fix(tests): make logs fetching more robust (#1878)

Co-authored-by: Heitor Lessa <lessa@amazon.com>

Author: rubenfonseca

tests/e2e/logger/test_logger.py

@@ -30,7 +30,7 @@ def test_basic_lambda_logs_visible(basic_handler_fn, basic_handler_fn_arn):
     data_fetcher.get_lambda_response(lambda_arn=basic_handler_fn_arn, payload=payload)
 
     # THEN
-    logs = data_fetcher.get_logs(function_name=basic_handler_fn, start_time=execution_time)
+    logs = data_fetcher.get_logs(function_name=basic_handler_fn, start_time=execution_time, minimum_log_entries=2)
 
     assert len(logs) == 2
     assert len(logs.get_cold_start_log()) == 1

tests/e2e/utils/data_fetcher/logs.py

@@ -29,6 +29,7 @@ def __init__(
         start_time: datetime,
         log_client: Optional[CloudWatchLogsClient] = None,
         filter_expression: Optional[str] = None,
+        minimum_log_entries: int = 1,
     ):
         """Fetch and expose Powertools Logger logs from CloudWatch Logs
 
@@ -42,12 +43,15 @@ def __init__(
             Amazon CloudWatch Logs Client, by default boto3.client('logs)
         filter_expression : Optional[str], optional
             CloudWatch Logs Filter Pattern expression, by default "message"
+        minimum_log_entries: int
+            Minimum number of log entries to be retrieved before exhausting retry attempts
         """
         self.function_name = function_name
         self.start_time = int(start_time.timestamp())
         self.log_client = log_client or boto3.client("logs")
         self.filter_expression = filter_expression or "message"  # Logger message key
         self.log_group = f"/aws/lambda/{self.function_name}"
+        self.minimum_log_entries = minimum_log_entries
         self.logs: List[Log] = self._get_logs()
 
     def get_log(self, key: str, value: Optional[any] = None) -> List[Log]:
@@ -112,6 +116,11 @@ def _get_logs(self) -> List[Log]:
                 continue
             filtered_logs.append(message)
 
+        if len(filtered_logs) < self.minimum_log_entries:
+            raise ValueError(
+                f"Number of log entries found doesn't meet minimum required ({self.minimum_log_entries}). Repeating..."
+            )
+
         return filtered_logs
 
     def __len__(self) -> int:
@@ -122,6 +131,7 @@ def __len__(self) -> int:
 def get_logs(
     function_name: str,
     start_time: datetime,
+    minimum_log_entries: int = 1,
     filter_expression: Optional[str] = None,
     log_client: Optional[CloudWatchLogsClient] = None,
 ) -> LogFetcher:
@@ -133,6 +143,8 @@ def get_logs(
         Name of Lambda function to fetch logs for
     start_time : datetime
         Start date range to filter traces
+    minimum_log_entries : int
+        Minimum number of log entries to be retrieved before exhausting retry attempts
     log_client : Optional[CloudWatchLogsClient], optional
         Amazon CloudWatch Logs Client, by default boto3.client('logs)
     filter_expression : Optional[str], optional
@@ -144,5 +156,9 @@ def get_logs(
         LogFetcher instance with logs available as properties and methods
     """
     return LogFetcher(
-        function_name=function_name, start_time=start_time, filter_expression=filter_expression, log_client=log_client
+        function_name=function_name,
+        start_time=start_time,
+        filter_expression=filter_expression,
+        log_client=log_client,
+        minimum_log_entries=minimum_log_entries,
     )