docs(data-classes): Add more cognito code examples (#340)

Author: Michael Brewer

Diff for: aws_lambda_powertools/utilities/data_classes/common.py

@@ -19,6 +19,11 @@ def __eq__(self, other: Any) -> bool:
     def get(self, key: str) -> Optional[Any]:
         return self._data.get(key)
 
+    @property
+    def raw_event(self) -> Dict[str, Any]:
+        """The original raw event dict"""
+        return self._data
+
 
 def get_header_value(headers: Dict[str, str], name: str, default_value: str, case_sensitive: bool) -> Optional[str]:
     """Get header value by name"""

Diff for: docs/utilities/data_classes.md

@@ -236,6 +236,8 @@ Define Auth Challenge | `data_classes.cognito_user_pool_event.DefineAuthChalleng
 Create Auth Challenge | `data_classes.cognito_user_pool_event.CreateAuthChallengeTriggerEvent`
 Verify Auth Challenge | `data_classes.cognito_user_pool_event.VerifyAuthChallengeResponseTriggerEvent`
 
+#### Post Confirmation Example
+
 === "app.py"
 
     ```python
@@ -248,6 +250,201 @@ Verify Auth Challenge | `data_classes.cognito_user_pool_event.VerifyAuthChalleng
         do_something_with(user_attributes)
     ```
 
+#### Define Auth Challenge Example
+
+!!! warning "NOTE "
+    In this example we are modifying the wrapped dict response fields, so we need to return the json serializable wrapped event in `event.raw_event`
+
+!!! info "NOTE "
+    This example is based on the AWS Cognito docs for [Define Auth Challenge Lambda Trigger](https://docs.aws.amazon.com/cognito/latest/developerguide/user-pool-lambda-define-auth-challenge.html){target="_blank"}
+
+=== "app.py"
+
+    ```python
+    from aws_lambda_powertools.utilities.data_classes.cognito_user_pool_event import DefineAuthChallengeTriggerEvent
+
+    def handler(event: dict, context) -> dict:
+        event: DefineAuthChallengeTriggerEvent = DefineAuthChallengeTriggerEvent(event)
+        if (
+            len(event.request.session) == 1
+            and event.request.session[0].challenge_name == "SRP_A"
+        ):
+            event.response.issue_tokens = False
+            event.response.fail_authentication = False
+            event.response.challenge_name = "PASSWORD_VERIFIER"
+        elif (
+            len(event.request.session) == 2
+            and event.request.session[1].challenge_name == "PASSWORD_VERIFIER"
+            and event.request.session[1].challenge_result
+        ):
+            event.response.issue_tokens = False
+            event.response.fail_authentication = False
+            event.response.challenge_name = "CUSTOM_CHALLENGE"
+        elif (
+            len(event.request.session) == 3
+            and event.request.session[2].challenge_name == "CUSTOM_CHALLENGE"
+            and event.request.session[2].challenge_result
+        ):
+            event.response.issue_tokens = True
+            event.response.fail_authentication = False
+        else:
+            event.response.issue_tokens = False
+            event.response.fail_authentication = True
+
+        return event.raw_event
+    ```
+=== "SPR_A response"
+
+    ```json hl_lines="25-27"
+    {
+        "version": "1",
+        "region": "us-east-1",
+        "userPoolId": "us-east-1_example",
+        "userName": "UserName",
+        "callerContext": {
+            "awsSdkVersion": "awsSdkVersion",
+            "clientId": "clientId"
+        },
+        "triggerSource": "DefineAuthChallenge_Authentication",
+        "request": {
+            "userAttributes": {
+                "sub": "4A709A36-7D63-4785-829D-4198EF10EBDA",
+                "email_verified": "true",
+                "name": "First Last",
+                "email": "define-auth@mail.com"
+            },
+            "session": [
+                {
+                    "challengeName": "SRP_A",
+                    "challengeResult": true
+                }
+            ]
+        },
+        "response": {
+            "issueTokens": false,
+            "failAuthentication": false,
+            "challengeName": "PASSWORD_VERIFIER"
+        }
+    }
+    ```
+=== "PASSWORD_VERIFIER success response"
+
+    ```json hl_lines="30-32"
+    {
+        "version": "1",
+        "region": "us-east-1",
+        "userPoolId": "us-east-1_example",
+        "userName": "UserName",
+        "callerContext": {
+            "awsSdkVersion": "awsSdkVersion",
+            "clientId": "clientId"
+        },
+        "triggerSource": "DefineAuthChallenge_Authentication",
+        "request": {
+            "userAttributes": {
+                "sub": "4A709A36-7D63-4785-829D-4198EF10EBDA",
+                "email_verified": "true",
+                "name": "First Last",
+                "email": "define-auth@mail.com"
+            },
+            "session": [
+                {
+                    "challengeName": "SRP_A",
+                    "challengeResult": true
+                },
+                {
+                    "challengeName": "PASSWORD_VERIFIER",
+                    "challengeResult": true
+                }
+            ]
+        },
+        "response": {
+            "issueTokens": false,
+            "failAuthentication": false,
+            "challengeName": "CUSTOM_CHALLENGE"
+        }
+    }
+
+    ```
+=== "CUSTOM_CHALLENGE success response"
+
+    ```json hl_lines="34 35"
+    {
+        "version": "1",
+        "region": "us-east-1",
+        "userPoolId": "us-east-1_example",
+        "userName": "UserName",
+        "callerContext": {
+            "awsSdkVersion": "awsSdkVersion",
+            "clientId": "clientId"
+        },
+        "triggerSource": "DefineAuthChallenge_Authentication",
+        "request": {
+            "userAttributes": {
+                "sub": "4A709A36-7D63-4785-829D-4198EF10EBDA",
+                "email_verified": "true",
+                "name": "First Last",
+                "email": "define-auth@mail.com"
+            },
+            "session": [
+                {
+                    "challengeName": "SRP_A",
+                    "challengeResult": true
+                },
+                {
+                    "challengeName": "PASSWORD_VERIFIER",
+                    "challengeResult": true
+                },
+                {
+                    "challengeName": "CUSTOM_CHALLENGE",
+                    "challengeResult": true
+                }
+            ]
+        },
+        "response": {
+            "issueTokens": true,
+            "failAuthentication": false
+        }
+    }
+    ```
+
+#### Create Auth Challenge Example
+
+!!! info "NOTE "
+    This example is based on the AWS Cognito docs for [Create Auth Challenge Lambda Trigger](https://docs.aws.amazon.com/cognito/latest/developerguide/user-pool-lambda-create-auth-challenge.html){target="_blank"}
+
+=== "app.py"
+
+    ```python
+    from aws_lambda_powertools.utilities.data_classes.cognito_user_pool_event import CreateAuthChallengeTriggerEvent
+
+    def handler(event: dict, context) -> dict:
+        event: CreateAuthChallengeTriggerEvent = CreateAuthChallengeTriggerEvent(event)
+        if event.request.challenge_name == "CUSTOM_CHALLENGE":
+            event.response.public_challenge_parameters = {"captchaUrl": "url/123.jpg"}
+            event.response.private_challenge_parameters = {"answer": "5"}
+            event.response.challenge_metadata = "CAPTCHA_CHALLENGE"
+        return event.raw_event
+    ```
+
+#### Verify Auth Challenge Response Example
+
+!!! info "NOTE "
+    This example is based on the AWS Cognito docs for [Verify Auth Challenge Response Lambda Trigger](https://docs.aws.amazon.com/cognito/latest/developerguide/user-pool-lambda-verify-auth-challenge-response.html){target="_blank"}
+
+=== "app.py"
+
+    ```python
+    from aws_lambda_powertools.utilities.data_classes.cognito_user_pool_event import VerifyAuthChallengeResponseTriggerEvent
+
+    def handler(event: dict, context) -> dict:
+        event: VerifyAuthChallengeResponseTriggerEvent = VerifyAuthChallengeResponseTriggerEvent(event)
+        event.response.answer_correct = (
+            event.request.private_challenge_parameters.get("answer") == event.request.challenge_answer
+        )
+        return event.raw_event
+    ```
+
 ### Connect Contact Flow
 
 === "app.py"

Diff for: tests/functional/test_lambda_trigger_events.py

@@ -73,6 +73,8 @@ def message(self) -> str:
     assert DataClassSample(data1) is not data1
     assert data1 is not DataClassSample(data1)
 
+    assert DataClassSample(data1).raw_event is data1
+
 
 def test_cloud_watch_trigger_event():
     event = CloudWatchLogsEvent(load_event("cloudWatchLogEvent.json"))
@@ -81,7 +83,7 @@ def test_cloud_watch_trigger_event():
     assert event.decompress_logs_data == decompressed_logs_data
 
     json_logs_data = event.parse_logs_data()
-    assert event.parse_logs_data()._data == json_logs_data._data
+    assert event.parse_logs_data().raw_event == json_logs_data.raw_event
     log_events = json_logs_data.log_events
     log_event = log_events[0]
 
@@ -97,7 +99,7 @@ def test_cloud_watch_trigger_event():
     assert log_event.extracted_fields is None
 
     event2 = CloudWatchLogsEvent(load_event("cloudWatchLogEvent.json"))
-    assert event._data == event2._data
+    assert event.raw_event == event2.raw_event
 
 
 def test_cognito_pre_signup_trigger_event():
@@ -493,7 +495,7 @@ def test_s3_trigger_event():
     assert s3.get_object.version_id is None
     assert s3.get_object.sequencer == "0C0F6F405D6ED209E1"
     assert record.glacier_event_data is None
-    assert event.record._data == event["Records"][0]
+    assert event.record.raw_event == event["Records"][0]
     assert event.bucket_name == "lambda-artifacts-deafc19498e3f2df"
     assert event.object_key == "b21b84d653bb07b05b1e6b33684dc11b"
 
@@ -559,7 +561,7 @@ def test_ses_trigger_event():
     assert headers[0].value == "<janedoe@example.com>"
     common_headers = mail.common_headers
     assert common_headers.return_path == "janedoe@example.com"
-    assert common_headers.get_from == common_headers._data["from"]
+    assert common_headers.get_from == common_headers.raw_event["from"]
     assert common_headers.date == "Wed, 7 Oct 2015 12:34:56 -0700"
     assert common_headers.to == [expected_address]
     assert common_headers.message_id == "<0123456789example.com>"
@@ -573,12 +575,12 @@ def test_ses_trigger_event():
     assert receipt.spf_verdict.status == "PASS"
     assert receipt.dmarc_verdict.status == "PASS"
     action = receipt.action
-    assert action.get_type == action._data["type"]
-    assert action.function_arn == action._data["functionArn"]
-    assert action.invocation_type == action._data["invocationType"]
-    assert event.record._data == event["Records"][0]
-    assert event.mail._data == event["Records"][0]["ses"]["mail"]
-    assert event.receipt._data == event["Records"][0]["ses"]["receipt"]
+    assert action.get_type == action.raw_event["type"]
+    assert action.function_arn == action.raw_event["functionArn"]
+    assert action.invocation_type == action.raw_event["invocationType"]
+    assert event.record.raw_event == event["Records"][0]
+    assert event.mail.raw_event == event["Records"][0]["ses"]["mail"]
+    assert event.receipt.raw_event == event["Records"][0]["ses"]["receipt"]
 
 
 def test_sns_trigger_event():
@@ -604,7 +606,7 @@ def test_sns_trigger_event():
     assert sns.unsubscribe_url == "https://sns.us-east-2.amazonaws.com/?Action=Unsubscribe"
     assert sns.topic_arn == "arn:aws:sns:us-east-2:123456789012:sns-lambda"
     assert sns.subject == "TestInvoke"
-    assert event.record._data == event["Records"][0]
+    assert event.record.raw_event == event["Records"][0]
     assert event.sns_message == "Hello from SNS!"