From d7f357cf5131b079d7cffe43c4b3951332819432 Mon Sep 17 00:00:00 2001 From: Michael Brewer Date: Wed, 8 Dec 2021 10:37:22 -0800 Subject: [PATCH] fix(event-sources): handle claimsOverrideDetails set to null --- aws_lambda_powertools/metrics/metric.py | 6 ++++-- .../utilities/data_classes/cognito_user_pool_event.py | 5 +++-- tests/functional/test_data_classes.py | 10 ++++++++++ 3 files changed, 17 insertions(+), 4 deletions(-) diff --git a/aws_lambda_powertools/metrics/metric.py b/aws_lambda_powertools/metrics/metric.py index 76ff4339dea..94b427738a1 100644 --- a/aws_lambda_powertools/metrics/metric.py +++ b/aws_lambda_powertools/metrics/metric.py @@ -1,7 +1,7 @@ import json import logging from contextlib import contextmanager -from typing import Dict, Optional, Union, Generator +from typing import Dict, Generator, Optional, Union from .base import MetricManager, MetricUnit @@ -61,7 +61,9 @@ def add_metric(self, name: str, unit: Union[MetricUnit, str], value: float) -> N @contextmanager -def single_metric(name: str, unit: MetricUnit, value: float, namespace: Optional[str] = None) -> Generator[SingleMetric, None, None]: +def single_metric( + name: str, unit: MetricUnit, value: float, namespace: Optional[str] = None +) -> Generator[SingleMetric, None, None]: """Context manager to simplify creation of a single metric Example diff --git a/aws_lambda_powertools/utilities/data_classes/cognito_user_pool_event.py b/aws_lambda_powertools/utilities/data_classes/cognito_user_pool_event.py index e467875305f..954d3d15b5f 100644 --- a/aws_lambda_powertools/utilities/data_classes/cognito_user_pool_event.py +++ b/aws_lambda_powertools/utilities/data_classes/cognito_user_pool_event.py @@ -523,8 +523,9 @@ def set_group_configuration_preferred_role(self, value: str): class PreTokenGenerationTriggerEventResponse(DictWrapper): @property def claims_override_details(self) -> ClaimsOverrideDetails: - # Ensure we have a `claimsOverrideDetails` element - self._data["response"].setdefault("claimsOverrideDetails", {}) + # Ensure we have a `claimsOverrideDetails` element and is not set to None + if self._data["response"].get("claimsOverrideDetails") is None: + self._data["response"]["claimsOverrideDetails"] = {} return ClaimsOverrideDetails(self._data["response"]["claimsOverrideDetails"]) diff --git a/tests/functional/test_data_classes.py b/tests/functional/test_data_classes.py index 5514a888e7d..ded32639233 100644 --- a/tests/functional/test_data_classes.py +++ b/tests/functional/test_data_classes.py @@ -281,6 +281,16 @@ def test_cognito_pre_token_generation_trigger_event(): assert claims_override_details.group_configuration.preferred_role == "role_name" assert event["response"]["claimsOverrideDetails"]["groupOverrideDetails"]["preferredRole"] == "role_name" + # Ensure that even if "claimsOverrideDetails" was explicitly set to None + # accessing `event.response.claims_override_details` would set it to `{}` + event["response"]["claimsOverrideDetails"] = None + claims_override_details = event.response.claims_override_details + assert claims_override_details._data == {} + assert event["response"]["claimsOverrideDetails"] == {} + claims_override_details.claims_to_suppress = ["email"] + assert claims_override_details.claims_to_suppress[0] == "email" + assert event["response"]["claimsOverrideDetails"]["claimsToSuppress"] == ["email"] + def test_cognito_define_auth_challenge_trigger_event(): event = DefineAuthChallengeTriggerEvent(load_event("cognitoDefineAuthChallengeEvent.json"))