Refactoring tests

Author: leandrodamascena

aws_lambda_powertools/utilities/data_masking/base.py

@@ -19,6 +19,7 @@
     DataMaskingUnsupportedTypeError,
 )
 from aws_lambda_powertools.utilities.data_masking.provider import BaseProvider
+from aws_lambda_powertools.warnings import PowertoolsUserWarning
 
 if TYPE_CHECKING:
     from numbers import Number
@@ -102,6 +103,7 @@ def erase(
         self,
         data: Any,
         fields: list[str] | None = None,
+        *,
         dynamic_mask: bool | None = None,
         custom_mask: str | None = None,
         regex_pattern: str | None = None,
@@ -311,11 +313,11 @@ def _apply_masking_rules(self, data: dict, masking_rules: dict) -> dict:
                             match.full_path.update(result, masked_value)
 
                     except Exception as e:
-                        warnings.warn(f"Error masking value for path {path}: {str(e)}", stacklevel=2)
+                        warnings.warn(f"Error masking value for path {path}: {str(e)}", category=PowertoolsUserWarning, stacklevel=2)
                         continue
 
             except Exception as e:
-                warnings.warn(f"Error processing path {path}: {str(e)}", stacklevel=2)
+                warnings.warn(f"Error processing path {path}: {str(e)}", category=PowertoolsUserWarning, stacklevel=2)
                 continue
 
         return result

tests/functional/data_masking/required_dependencies/test_erase_data_masking.py

@@ -8,6 +8,7 @@
     DataMaskingFieldNotFoundError,
     DataMaskingUnsupportedTypeError,
 )
+from aws_lambda_powertools.warnings import PowertoolsUserWarning
 
 
 @pytest.fixture
@@ -304,7 +305,7 @@ def erase(self, value, **kwargs):
     data_masker.provider = MockProvider()
 
     # WHEN erase is called
-    with pytest.warns(UserWarning, match="Error masking value for path value: Mock error"):
+    with pytest.warns(expected_warning=PowertoolsUserWarning, match="Error masking value for path value: Mock error"):
         masked_data = data_masker.erase(data, masking_rules={"value": {"rule": "value"}})
 
     # THEN the original data should remain unchanged
@@ -422,5 +423,5 @@ def test_erase_dictionary_with_masking_rules_wrong_field(data_masker):
 
     # WHEN erase is called with wrong masking rules
     # We must have a warning
-    with pytest.warns(UserWarning, match="Error processing path*"):
+    with pytest.warns(expected_warning=PowertoolsUserWarning, match="Error processing path*"):
         data_masker.erase(data, masking_rules=masking_rules)