refactor(tests): add LambdaContext type in tests (#6214)

basvandriel · leandrodamascena · web-flow · commit 83222ddd5010 · 2025-03-07T17:34:50.000Z
* Type `LambdaContext` types in tests

* Remove types on parameters

* Fix highlight

---------

Co-authored-by: Leandro Damascena &lt;lcdama@amazon.pt&gt;
diff --git a/docs/core/event_handler/appsync.md b/docs/core/event_handler/appsync.md
@@ -586,7 +586,7 @@ Here's an example of how you can test your synchronous resolvers:
 
 === "assert_graphql_response_module.py"
 
-    ```python hl_lines="11"
+    ```python hl_lines="10"
     --8<-- "examples/event_handler_graphql/src/assert_graphql_response_module.py"
     ```
 
diff --git a/docs/core/event_handler/bedrock_agents.md b/docs/core/event_handler/bedrock_agents.md
@@ -319,7 +319,7 @@ Test your routes by passing an [Agent for Amazon Bedrock proxy event](https://do
 
 === "assert_bedrock_agent_response.py"
 
-	```python hl_lines="21-23 27"
+	```python hl_lines="22-24 28"
 	--8<-- "examples/event_handler_bedrock_agents/src/assert_bedrock_agent_response.py"
 	```
 
diff --git a/docs/core/metrics.md b/docs/core/metrics.md
@@ -425,7 +425,7 @@ You can read standard output and assert whether metrics have been flushed. Here'
 
     This will be needed when using `capture_cold_start_metric=True`, or when both `Metrics` and `single_metric` are used.
 
-    ```python hl_lines="20-21 27"
+    ```python hl_lines="21-22 28"
     --8<-- "examples/metrics/src/assert_multiple_emf_blobs.py"
     ```
 
diff --git a/docs/utilities/data_masking.md b/docs/utilities/data_masking.md
@@ -671,7 +671,7 @@ Testing your code with a simple erase operation
 
 === "test_lambda_mask.py"
 
-```python hl_lines="22"
+```python
 --8<-- "examples/data_masking/tests/test_lambda_mask.py"
 ```
 
diff --git a/examples/batch_processing/src/getting_started_with_test.py b/examples/batch_processing/src/getting_started_with_test.py
@@ -11,15 +11,16 @@ def load_event(path: Path):
         return json.load(f)
 
 
-@pytest.fixture
-def lambda_context():
-    @dataclass
-    class LambdaContext:
-        function_name: str = "test"
-        memory_limit_in_mb: int = 128
-        invoked_function_arn: str = "arn:aws:lambda:eu-west-1:809313241:function:test"
-        aws_request_id: str = "52fdfc07-2182-154f-163f-5f0f9a621d72"
+@dataclass
+class LambdaContext:
+    function_name: str = "test"
+    memory_limit_in_mb: int = 128
+    invoked_function_arn: str = "arn:aws:lambda:eu-west-1:809313241:function:test"
+    aws_request_id: str = "52fdfc07-2182-154f-163f-5f0f9a621d72"
+
 
+@pytest.fixture
+def lambda_context() -> LambdaContext:
     return LambdaContext()
 
 
@@ -29,7 +30,7 @@ def sqs_event():
     return load_event(path=Path("events/sqs_event.json"))
 
 
-def test_app_batch_partial_response(sqs_event, lambda_context):
+def test_app_batch_partial_response(sqs_event, lambda_context: LambdaContext):
     # GIVEN
     processor_result = processor  # access processor for additional assertions
     successful_record = sqs_event["Records"][0]
diff --git a/examples/data_masking/tests/test_lambda_mask.py b/examples/data_masking/tests/test_lambda_mask.py
@@ -4,18 +4,19 @@
 import test_lambda_mask
 
 
-@pytest.fixture
-def lambda_context():
-    @dataclass
-    class LambdaContext:
-        function_name: str = "test"
-        memory_limit_in_mb: int = 128
-        invoked_function_arn: str = "arn:aws:lambda:eu-west-1:111111111:function:test"
-        aws_request_id: str = "52fdfc07-2182-154f-163f-5f0f9a621d72"
+@dataclass
+class LambdaContext:
+    function_name: str = "test"
+    memory_limit_in_mb: int = 128
+    invoked_function_arn: str = "arn:aws:lambda:eu-west-1:111111111:function:test"
+    aws_request_id: str = "52fdfc07-2182-154f-163f-5f0f9a621d72"
+
+    def get_remaining_time_in_millis(self) -> int:
+        return 5
 
-        def get_remaining_time_in_millis(self) -> int:
-            return 5
 
+@pytest.fixture
+def lambda_context() -> LambdaContext:
     return LambdaContext()
 
 
diff --git a/examples/event_handler_bedrock_agents/src/assert_bedrock_agent_response.py b/examples/event_handler_bedrock_agents/src/assert_bedrock_agent_response.py
@@ -4,19 +4,20 @@
 import pytest
 
 
-@pytest.fixture
-def lambda_context():
-    @dataclass
-    class LambdaContext:
-        function_name: str = "test"
-        memory_limit_in_mb: int = 128
-        invoked_function_arn: str = "arn:aws:lambda:eu-west-1:123456789012:function:test"
-        aws_request_id: str = "da658bd3-2d6f-4e7b-8ec2-937234644fdc"
+@dataclass
+class LambdaContext:
+    function_name: str = "test"
+    memory_limit_in_mb: int = 128
+    invoked_function_arn: str = "arn:aws:lambda:eu-west-1:123456789012:function:test"
+    aws_request_id: str = "da658bd3-2d6f-4e7b-8ec2-937234644fdc"
+
 
+@pytest.fixture
+def lambda_context() -> LambdaContext:
     return LambdaContext()
 
 
-def test_lambda_handler(lambda_context):
+def test_lambda_handler(lambda_context: LambdaContext):
     minimal_event = {
         "apiPath": "/current_time",
         "httpMethod": "GET",
diff --git a/examples/event_handler_graphql/src/assert_async_graphql_response.py b/examples/event_handler_graphql/src/assert_async_graphql_response.py
@@ -10,15 +10,16 @@
 )
 
 
-@pytest.fixture
-def lambda_context():
-    @dataclass
-    class LambdaContext:
-        function_name: str = "test"
-        memory_limit_in_mb: int = 128
-        invoked_function_arn: str = "arn:aws:lambda:eu-west-1:123456789012:function:test"
-        aws_request_id: str = "da658bd3-2d6f-4e7b-8ec2-937234644fdc"
+@dataclass
+class LambdaContext:
+    function_name: str = "test"
+    memory_limit_in_mb: int = 128
+    invoked_function_arn: str = "arn:aws:lambda:eu-west-1:123456789012:function:test"
+    aws_request_id: str = "da658bd3-2d6f-4e7b-8ec2-937234644fdc"
+
 
+@pytest.fixture
+def lambda_context() -> LambdaContext:
     return LambdaContext()
 
 
diff --git a/examples/event_handler_graphql/src/assert_graphql_response.py b/examples/event_handler_graphql/src/assert_graphql_response.py
@@ -8,15 +8,16 @@
 from assert_graphql_response_module import Location, app  # instance of AppSyncResolver
 
 
-@pytest.fixture
-def lambda_context():
-    @dataclass
-    class LambdaContext:
-        function_name: str = "test"
-        memory_limit_in_mb: int = 128
-        invoked_function_arn: str = "arn:aws:lambda:eu-west-1:123456789012:function:test"
-        aws_request_id: str = "da658bd3-2d6f-4e7b-8ec2-937234644fdc"
+@dataclass
+class LambdaContext:
+    function_name: str = "test"
+    memory_limit_in_mb: int = 128
+    invoked_function_arn: str = "arn:aws:lambda:eu-west-1:123456789012:function:test"
+    aws_request_id: str = "da658bd3-2d6f-4e7b-8ec2-937234644fdc"
+
 
+@pytest.fixture
+def lambda_context() -> LambdaContext:
     return LambdaContext()
 
 
diff --git a/examples/event_handler_rest/src/assert_alb_api_resolver_response.py b/examples/event_handler_rest/src/assert_alb_api_resolver_response.py
@@ -4,19 +4,20 @@
 import pytest
 
 
-@pytest.fixture
-def lambda_context():
-    @dataclass
-    class LambdaContext:
-        function_name: str = "test"
-        memory_limit_in_mb: int = 128
-        invoked_function_arn: str = "arn:aws:lambda:eu-west-1:123456789012:function:test"
-        aws_request_id: str = "da658bd3-2d6f-4e7b-8ec2-937234644fdc"
+@dataclass
+class LambdaContext:
+    function_name: str = "test"
+    memory_limit_in_mb: int = 128
+    invoked_function_arn: str = "arn:aws:lambda:eu-west-1:123456789012:function:test"
+    aws_request_id: str = "da658bd3-2d6f-4e7b-8ec2-937234644fdc"
+
 
+@pytest.fixture
+def lambda_context() -> LambdaContext:
     return LambdaContext()
 
 
-def test_lambda_handler(lambda_context):
+def test_lambda_handler(lambda_context: LambdaContext):
     minimal_event = {
         "path": "/todos",
         "httpMethod": "GET",
diff --git a/examples/event_handler_rest/src/assert_function_url_api_resolver_response.py b/examples/event_handler_rest/src/assert_function_url_api_resolver_response.py
@@ -4,19 +4,20 @@
 import pytest
 
 
-@pytest.fixture
-def lambda_context():
-    @dataclass
-    class LambdaContext:
-        function_name: str = "test"
-        memory_limit_in_mb: int = 128
-        invoked_function_arn: str = "arn:aws:lambda:eu-west-1:123456789012:function:test"
-        aws_request_id: str = "da658bd3-2d6f-4e7b-8ec2-937234644fdc"
+@dataclass
+class LambdaContext:
+    function_name: str = "test"
+    memory_limit_in_mb: int = 128
+    invoked_function_arn: str = "arn:aws:lambda:eu-west-1:123456789012:function:test"
+    aws_request_id: str = "da658bd3-2d6f-4e7b-8ec2-937234644fdc"
+
 
+@pytest.fixture
+def lambda_context() -> LambdaContext:
     return LambdaContext()
 
 
-def test_lambda_handler(lambda_context):
+def test_lambda_handler(lambda_context: LambdaContext):
     minimal_event = {
         "rawPath": "/todos",
         "requestContext": {
diff --git a/examples/event_handler_rest/src/assert_http_api_resolver_response.py b/examples/event_handler_rest/src/assert_http_api_resolver_response.py
@@ -4,19 +4,20 @@
 import pytest
 
 
-@pytest.fixture
-def lambda_context():
-    @dataclass
-    class LambdaContext:
-        function_name: str = "test"
-        memory_limit_in_mb: int = 128
-        invoked_function_arn: str = "arn:aws:lambda:eu-west-1:123456789012:function:test"
-        aws_request_id: str = "da658bd3-2d6f-4e7b-8ec2-937234644fdc"
+@dataclass
+class LambdaContext:
+    function_name: str = "test"
+    memory_limit_in_mb: int = 128
+    invoked_function_arn: str = "arn:aws:lambda:eu-west-1:123456789012:function:test"
+    aws_request_id: str = "da658bd3-2d6f-4e7b-8ec2-937234644fdc"
+
 
+@pytest.fixture
+def lambda_context() -> LambdaContext:
     return LambdaContext()
 
 
-def test_lambda_handler(lambda_context):
+def test_lambda_handler(lambda_context: LambdaContext):
     minimal_event = {
         "rawPath": "/todos",
         "requestContext": {
diff --git a/examples/event_handler_rest/src/assert_rest_api_resolver_response.py b/examples/event_handler_rest/src/assert_rest_api_resolver_response.py
@@ -4,15 +4,16 @@
 import pytest
 
 
-@pytest.fixture
-def lambda_context():
-    @dataclass
-    class LambdaContext:
-        function_name: str = "test"
-        memory_limit_in_mb: int = 128
-        invoked_function_arn: str = "arn:aws:lambda:eu-west-1:123456789012:function:test"
-        aws_request_id: str = "da658bd3-2d6f-4e7b-8ec2-937234644fdc"
+@dataclass
+class LambdaContext:
+    function_name: str = "test"
+    memory_limit_in_mb: int = 128
+    invoked_function_arn: str = "arn:aws:lambda:eu-west-1:123456789012:function:test"
+    aws_request_id: str = "da658bd3-2d6f-4e7b-8ec2-937234644fdc"
+
 
+@pytest.fixture
+def lambda_context() -> LambdaContext:
     return LambdaContext()
 
 
diff --git a/examples/idempotency/tests/test_disabling_idempotency_utility.py b/examples/idempotency/tests/test_disabling_idempotency_utility.py
@@ -4,22 +4,23 @@
 import pytest
 
 
-@pytest.fixture
-def lambda_context():
-    @dataclass
-    class LambdaContext:
-        function_name: str = "test"
-        memory_limit_in_mb: int = 128
-        invoked_function_arn: str = "arn:aws:lambda:eu-west-1:809313241:function:test"
-        aws_request_id: str = "52fdfc07-2182-154f-163f-5f0f9a621d72"
+@dataclass
+class LambdaContext:
+    function_name: str = "test"
+    memory_limit_in_mb: int = 128
+    invoked_function_arn: str = "arn:aws:lambda:eu-west-1:809313241:function:test"
+    aws_request_id: str = "52fdfc07-2182-154f-163f-5f0f9a621d72"
+
+    def get_remaining_time_in_millis(self) -> int:
+        return 5
 
-        def get_remaining_time_in_millis(self) -> int:
-            return 5
 
+@pytest.fixture
+def lambda_context() -> LambdaContext:
     return LambdaContext()
 
 
-def test_idempotent_lambda_handler(monkeypatch, lambda_context):
+def test_idempotent_lambda_handler(monkeypatch, lambda_context: LambdaContext):
     # Set POWERTOOLS_IDEMPOTENCY_DISABLED before calling decorated functions
     monkeypatch.setenv("POWERTOOLS_IDEMPOTENCY_DISABLED", 1)
 
diff --git a/examples/idempotency/tests/test_with_dynamodb_local.py b/examples/idempotency/tests/test_with_dynamodb_local.py
@@ -5,18 +5,19 @@
 import pytest
 
 
-@pytest.fixture
-def lambda_context():
-    @dataclass
-    class LambdaContext:
-        function_name: str = "test"
-        memory_limit_in_mb: int = 128
-        invoked_function_arn: str = "arn:aws:lambda:eu-west-1:809313241:function:test"
-        aws_request_id: str = "52fdfc07-2182-154f-163f-5f0f9a621d72"
+@dataclass
+class LambdaContext:
+    function_name: str = "test"
+    memory_limit_in_mb: int = 128
+    invoked_function_arn: str = "arn:aws:lambda:eu-west-1:809313241:function:test"
+    aws_request_id: str = "52fdfc07-2182-154f-163f-5f0f9a621d72"
+
+    def get_remaining_time_in_millis(self) -> int:
+        return 5
 
-        def get_remaining_time_in_millis(self) -> int:
-            return 5
 
+@pytest.fixture
+def lambda_context() -> LambdaContext:
     return LambdaContext()
 
 
diff --git a/examples/idempotency/tests/test_with_io_operations.py b/examples/idempotency/tests/test_with_io_operations.py
@@ -5,18 +5,19 @@
 import pytest
 
 
-@pytest.fixture
-def lambda_context():
-    @dataclass
-    class LambdaContext:
-        function_name: str = "test"
-        memory_limit_in_mb: int = 128
-        invoked_function_arn: str = "arn:aws:lambda:eu-west-1:809313241:function:test"
-        aws_request_id: str = "52fdfc07-2182-154f-163f-5f0f9a621d72"
+@dataclass
+class LambdaContext:
+    function_name: str = "test"
+    memory_limit_in_mb: int = 128
+    invoked_function_arn: str = "arn:aws:lambda:eu-west-1:809313241:function:test"
+    aws_request_id: str = "52fdfc07-2182-154f-163f-5f0f9a621d72"
+
+    def get_remaining_time_in_millis(self) -> int:
+        return 5
 
-        def get_remaining_time_in_millis(self) -> int:
-            return 5
 
+@pytest.fixture
+def lambda_context() -> LambdaContext:
     return LambdaContext()
 
 
diff --git a/examples/logger/src/fake_lambda_context_for_logger.py b/examples/logger/src/fake_lambda_context_for_logger.py
diff --git a/examples/metrics/src/assert_multiple_emf_blobs.py b/examples/metrics/src/assert_multiple_emf_blobs.py
diff --git a/tests/functional/idempotency/_boto3/conftest.py b/tests/functional/idempotency/_boto3/conftest.py
diff --git a/tests/functional/idempotency/_redis/test_redis_layer.py b/tests/functional/idempotency/_redis/test_redis_layer.py
diff --git a/tests/integration/idempotency/test_idempotency_redis.py b/tests/integration/idempotency/test_idempotency_redis.py
