refactor: create LogFetcher to simply assertion

Author: heitorlessa

aws_lambda_powertools/shared/constants.py

@@ -23,3 +23,12 @@
 XRAY_SDK_CORE_MODULE: str = "aws_xray_sdk.core"
 
 IDEMPOTENCY_DISABLED_ENV: str = "POWERTOOLS_IDEMPOTENCY_DISABLED"
+
+LOGGER_LAMBDA_CONTEXT_KEYS = [
+    "function_arn",
+    "function_memory_size",
+    "function_name",
+    "function_request_id",
+    "cold_start",
+    "xray_trace_id",
+]

tests/e2e/logger/test_logger.py

@@ -18,22 +18,17 @@ def basic_handler_fn_arn(infrastructure: dict) -> str:
 
 def test_basic_lambda_logs_visible(basic_handler_fn, basic_handler_fn_arn):
     # GIVEN
-    required_keys = (
-        "xray_trace_id",
-        "function_request_id",
-        "function_arn",
-        "function_memory_size",
-        "function_name",
-        "cold_start",
-    )
     message = "logs should be visible with default settings"
-    additional_keys = {"order_id": f"{uuid4()}"}
+    custom_key = "order_id"
+    additional_keys = {custom_key: f"{uuid4()}"}
     payload = json.dumps({"message": message, "append_keys": additional_keys})
 
     # WHEN
     _, execution_time = data_fetcher.get_lambda_response(lambda_arn=basic_handler_fn_arn, payload=payload)
-    filtered_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)
 
     # THEN
-    assert all(keys in logs.dict(exclude_unset=True) for logs in filtered_logs for keys in required_keys)
-    assert any(getattr(logs, "order_id", False) for logs in filtered_logs)
+    assert len(logs) == 1
+    assert len(logs.get_cold_start_log()) == 1
+    assert len(logs.get_log(key=custom_key)) == 1
+    assert logs.have_logger_context_keys() is True

tests/e2e/utils/data_fetcher/logs.py

@@ -7,6 +7,8 @@
 from pydantic import BaseModel, Extra
 from retry import retry
 
+from aws_lambda_powertools.shared.constants import LOGGER_LAMBDA_CONTEXT_KEYS
+
 
 class Log(BaseModel, extra=Extra.allow):
     level: str
@@ -22,31 +24,68 @@ class Log(BaseModel, extra=Extra.allow):
     xray_trace_id: Optional[str]
 
 
+class LogFetcher:
+    def __init__(
+        self,
+        function_name: str,
+        start_time: datetime,
+        log_client: Optional[CloudWatchLogsClient] = None,
+        filter_expression: Optional[str] = None,
+    ) -> None:
+        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.logs: List[Log] = self.get_logs()
+
+    def get_logs(self):
+        ret = self.log_client.filter_log_events(
+            logGroupName=self.log_group,
+            startTime=self.start_time,
+            filterPattern=self.filter_expression,
+        )
+
+        if not ret["events"]:
+            raise ValueError("Empty response from Cloudwatch Logs. Repeating...")
+
+        filtered_logs = []
+        for event in ret["events"]:
+            try:
+                message = Log(**json.loads(event["message"]))
+            except json.decoder.JSONDecodeError:
+                continue
+            filtered_logs.append(message)
+
+        return filtered_logs
+
+    def get_log(self, key: str, value: Optional[any] = None) -> List[Log]:
+        logs = []
+        for log in self.logs:
+            log_value = getattr(log, key, None)
+            if value is not None and log_value == value:
+                logs.append(log)
+            if value is None and getattr(log, key, False):
+                logs.append(log)
+        return logs
+
+    def get_cold_start_log(self) -> List[Log]:
+        return [log for log in self.logs if log.cold_start]
+
+    def have_logger_context_keys(self) -> bool:
+        return all(getattr(log, key, False) for log in self.logs for key in LOGGER_LAMBDA_CONTEXT_KEYS)
+
+    def __len__(self) -> int:
+        return len(self.logs)
+
+
 @retry(ValueError, delay=2, jitter=1.5, tries=10)
 def get_logs(
     function_name: str,
     start_time: datetime,
-    log_client: Optional[CloudWatchLogsClient] = None,
     filter_expression: Optional[str] = None,
-) -> List[Log]:
-    log_client = log_client or boto3.client("logs")
-    filter_expression = filter_expression or "message"  # Logger message key
-
-    response = log_client.filter_log_events(
-        logGroupName=f"/aws/lambda/{function_name}",
-        startTime=int(start_time.timestamp()),
-        filterPattern=filter_expression,
+    log_client: Optional[CloudWatchLogsClient] = None,
+) -> LogFetcher:
+    return LogFetcher(
+        function_name=function_name, start_time=start_time, filter_expression=filter_expression, log_client=log_client
     )
-
-    if not response["events"]:
-        raise ValueError("Empty response from Cloudwatch Logs. Repeating...")
-
-    filtered_logs = []
-    for event in response["events"]:
-        try:
-            message = Log(**json.loads(event["message"]))
-        except json.decoder.JSONDecodeError:
-            continue
-        filtered_logs.append(message)
-
-    return filtered_logs