feat: make log_client optional and fix client type

Author: heitorlessa

poetry.lock

@@ -65,6 +65,7 @@ mypy-boto3-lambda = "^1.24.0"
 mypy-boto3-xray = "^1.24.0"
 mypy-boto3-s3 = { version = "^1.24.0", python = ">=3.7" }
 mypy-boto3-cloudformation = { version = "^1.24.0", python = ">=3.7" }
+mypy-boto3-logs = { version = "^1.24.0", python = ">=3.7" }
 types-requests = "^2.28.8"
 typing-extensions = { version = "^4.3.0", python = ">=3.7" }
 python-snappy = "^0.6.1"

pyproject.toml

@@ -1,4 +1,3 @@
-import boto3
 import pytest
 
 from tests.e2e.utils import data_fetcher
@@ -16,7 +15,6 @@ def basic_handler_fn_arn(infrastructure: dict) -> str:
 
 def test_basic_lambda_logs_visible(basic_handler_fn, basic_handler_fn_arn):
     # GIVEN
-    cw_client = boto3.client("logs")
     required_keys = (
         "xray_trace_id",
         "function_request_id",
@@ -28,9 +26,7 @@ def test_basic_lambda_logs_visible(basic_handler_fn, basic_handler_fn_arn):
 
     # WHEN
     _, execution_time = data_fetcher.get_lambda_response(lambda_arn=basic_handler_fn_arn)
-    filtered_logs = data_fetcher.get_logs(
-        lambda_function_name=basic_handler_fn, start_time=execution_time, log_client=cw_client
-    )
+    filtered_logs = data_fetcher.get_logs(lambda_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)

tests/e2e/logger/test_logger.py

@@ -3,7 +3,8 @@
 from functools import lru_cache
 from typing import List, Optional, Union
 
-from mypy_boto3_cloudwatch import CloudWatchClient
+import boto3
+from mypy_boto3_logs import CloudWatchLogsClient
 from pydantic import BaseModel
 from retry import retry
 
@@ -25,10 +26,15 @@ class Log(BaseModel):
 
 @lru_cache(maxsize=10, typed=False)
 @retry(ValueError, delay=1, jitter=1, tries=20)
-def get_logs(lambda_function_name: str, log_client: CloudWatchClient, start_time: datetime) -> List[Log]:
+def get_logs(
+    lambda_function_name: str, start_time: datetime, log_client: Optional[CloudWatchLogsClient] = None
+) -> List[Log]:
+    log_client = log_client or boto3.client("logs")
+
     response = log_client.filter_log_events(
         logGroupName=f"/aws/lambda/{lambda_function_name}", startTime=int(start_time.timestamp())
     )
+
     if not response["events"]:
         raise ValueError("Empty response from Cloudwatch Logs. Repeating...")