Small fix

seshubaws · Dec 18, 2023 · 74fccd7 · 74fccd7
1 parent fec33a6
commit 74fccd7
Show file tree

Hide file tree

Showing 3 changed files with 9 additions and 6 deletions.
diff --git a/aws_lambda_powertools/utilities/_data_masking/base.py b/aws_lambda_powertools/utilities/_data_masking/base.py
@@ -5,7 +5,7 @@
 from typing import Any, Callable, Iterable, Optional, Union
 
 from aws_lambda_powertools.utilities._data_masking.exceptions import (
-    DataMaskingFieldNotFound,
+    DataMaskingFieldNotFoundError,
     DataMaskingUnsupportedTypeError,
 )
 from aws_lambda_powertools.utilities._data_masking.provider import BaseProvider
@@ -180,7 +180,7 @@ def _apply_action_to_fields(
                         current_dict = current_dict[int(key)]
                 except KeyError:
                     # Handle the case when the key doesn't exist
-                    raise DataMaskingFieldNotFound(f"Key {key} not found in {current_dict}")
+                    raise DataMaskingFieldNotFoundError(f"Key {key} not found in {current_dict}")
 
             last_key = keys[-1]
 
@@ -225,6 +225,6 @@ def _apply_action_to_specific_type(self, current_dict: dict, action: Callable, l
             current_dict = set(elements_list)
         else:
             # Handle the case when the last key doesn't exist
-            raise DataMaskingFieldNotFound(f"Key {last_key} not found in {current_dict}")
+            raise DataMaskingFieldNotFoundError(f"Key {last_key} not found in {current_dict}")
 
         return current_dict
diff --git a/aws_lambda_powertools/utilities/_data_masking/exceptions.py b/aws_lambda_powertools/utilities/_data_masking/exceptions.py
@@ -28,7 +28,7 @@ class DataMaskingContextMismatchError(Exception):
     """
 
 
-class DataMaskingFieldNotFound(Exception):
+class DataMaskingFieldNotFoundError(Exception):
     """
     Field not found.
     """
diff --git a/tests/unit/data_masking/test_unit_data_masking.py b/tests/unit/data_masking/test_unit_data_masking.py
@@ -4,7 +4,10 @@
 
 from aws_lambda_powertools.utilities._data_masking.base import DataMasking
 from aws_lambda_powertools.utilities._data_masking.constants import DATA_MASKING_STRING
-from aws_lambda_powertools.utilities._data_masking.exceptions import DataMaskingUnsupportedTypeError
+from aws_lambda_powertools.utilities._data_masking.exceptions import (
+    DataMaskingFieldNotFoundError,
+    DataMaskingUnsupportedTypeError,
+)
 
 
 @pytest.fixture
@@ -169,7 +172,7 @@ def test_parsing_nonexistent_fields(data_masker):
     }
 
     # WHEN attempting to pass in fields that do not exist in the input data
-    with pytest.raises(KeyError):
+    with pytest.raises(DataMaskingFieldNotFoundError):
         # THEN the result is a KeyError
         data_masker.mask(data, ["3.1.True"])