add tests for new masks

Author: anafalcao

aws_lambda_powertools/utilities/data_masking/base.py

@@ -119,8 +119,6 @@ def erase(
         mask_format: str | None = None,
         masking_rules: dict | None = None,
     ) -> str | list[str] | tuple[str] | dict:
-        if not data:
-            return data
         if masking_rules:
             return self._apply_masking_rules(data, masking_rules)
         else:

aws_lambda_powertools/utilities/data_masking/provider/base.py

@@ -77,11 +77,16 @@ def erase(
         masking_rules: dict | None = None,
         **kwargs,
     ) -> Any:
-        # Handle empty or None data
-        if data is None or (isinstance(data, (str, list, dict)) and not data):
-            return data
 
-        result = data  # Default to returning the original data
+        result = DATA_MASKING_STRING
+
+        if not any([dynamic_mask, custom_mask, regex_pattern, mask_format, masking_rules]):
+            if isinstance(data, (str, int, float, dict, bytes)):
+                return DATA_MASKING_STRING
+            elif isinstance(data, (list, tuple, set)):
+                return type(data)([DATA_MASKING_STRING] * len(data))
+            else:
+                return DATA_MASKING_STRING
 
         if isinstance(data, (str, int, float)):
             result = self._mask_primitive(str(data), dynamic_mask, custom_mask, regex_pattern, mask_format, **kwargs)
@@ -194,7 +199,7 @@ def _regex_mask(self, data: str, regex_pattern: str, mask_format: str) -> str:
                 _regex_cache[regex_pattern] = re.compile(regex_pattern)
             return _regex_cache[regex_pattern].sub(mask_format, data)
         except re.error:
-            return DATA_MASKING_STRING
+            return data
 
     def _custom_erase(self, data: str, **kwargs) -> str:
         if not data:

tests/unit/data_masking/_aws_encryption_sdk/test_unit_data_masking.py

@@ -25,6 +25,16 @@ def test_erase_int(data_masker):
     assert erased_string == DATA_MASKING_STRING
 
 
+def test_erase_int_custom_mask(data_masker):
+    # GIVEN an int data type
+
+    # WHEN erase is called with no fields argument
+    erased_string = data_masker.erase(42, custom_mask="XX")
+
+    # THEN the result is the data masked
+    assert erased_string == "XX"
+
+
 def test_erase_float(data_masker):
     # GIVEN a float data type
 
@@ -205,3 +215,36 @@ def test_parsing_nonexistent_fields_warning_on_missing_field():
 
     # THEN the "erased" payload is the same of the original
     assert masked_json_string == data
+
+
+def test_regex_mask(data_masker):
+    data = "Hello! My name is Fulano Ciclano"
+    regex_pattern = r"\b[A-Z][a-z]+ [A-Z][a-z]+\b"
+    mask_format = "XXXX XXXX"
+
+    result = data_masker.erase(data, regex_pattern=regex_pattern, mask_format=mask_format)
+
+    assert result == "Hello! My name is XXXX XXXX"
+
+
+def test_erase_json_dict_with_fields_and_masks(data_masker):
+    # GIVEN the data type is a json representation of a dictionary
+    data = json.dumps(
+        {
+            "a": {
+                "1": {"None": "hello", "four": "world"},
+                "b": {"3": {"4": "goodbye", "e": "world"}},
+            },
+        },
+    )
+
+    # WHEN erase is called with a list of fields specified
+    masked_json_string = data_masker.erase(data, fields=["a.'1'.None", "a..'4'"], dynamic_mask=True)
+
+    # THEN the result is only the specified fields are erased
+    assert masked_json_string == {
+        "a": {
+            "1": {"None": "*****", "four": "world"},
+            "b": {"3": {"4": "*******", "e": "world"}},
+        },
+    }