Fixed docstrings, added a test

seshubaws · Jan 23, 2024 · 85766bf · 85766bf
1 parent 429eb8a
commit 85766bf
Show file tree

Hide file tree

Showing 3 changed files with 29 additions and 3 deletions.
diff --git a/aws_lambda_powertools/utilities/data_masking/base.py b/aws_lambda_powertools/utilities/data_masking/base.py
@@ -194,8 +194,10 @@ def _apply_action_to_fields(
                 value of the field as the first argument and any additional arguments that might be required
                 for the action. It performs an operation on the current value using the provided arguments and
                 returns the modified value.
-            **provider_options:
-                Additional keyword arguments to pass to the 'action' function.
+            provider_options : dict
+                Optional dictionary representing additional options for the action.
+            **encryption_context: str
+                Additional keyword arguments collected into a dictionary.
 
         Returns
         -------

diff --git a/aws_lambda_powertools/utilities/data_masking/provider/kms/aws_encryption_sdk.py b/aws_lambda_powertools/utilities/data_masking/provider/kms/aws_encryption_sdk.py
@@ -146,8 +146,10 @@ def encrypt(self, data: Any, provider_options: dict | None = None, **encryption_
         -------
             data : Union[bytes, str]
                 The data to be encrypted.
-            provider_options
+            provider_options : dict
                 Additional options for the aws_encryption_sdk.EncryptionSDKClient
+            **encryption_context : str
+                Additional keyword arguments collected into a dictionary.
 
         Returns
         -------

diff --git a/tests/functional/data_masking/test_aws_encryption_sdk.py b/tests/functional/data_masking/test_aws_encryption_sdk.py
@@ -7,6 +7,8 @@
 
 import pytest
 
+from aws_encryption_sdk.identifiers import Algorithm
+
 from aws_lambda_powertools.utilities.data_masking import DataMasking
 from aws_lambda_powertools.utilities.data_masking.constants import DATA_MASKING_STRING
 from aws_lambda_powertools.utilities.data_masking.provider import BaseProvider
@@ -459,3 +461,23 @@ def test_encrypt_with_complex_search(data_masker):
 
     # THEN the result is only the specified fields are masked
     assert decrypted_data == json.loads(data)
+
+def test_encrypt_with_provider_options(data_masker):
+    # GIVEN the data type is a json representation of a dictionary with a list inside
+    data = json.dumps(
+        {
+            "payload": {
+                "first": ["value1", "value2"],
+                "second": (0, 1),
+            },
+        },
+    )
+
+    fields_operation = ["payload.first[0]", "payload.second[0]"]
+    provider_options = {"algorithm": Algorithm.AES_256_GCM_HKDF_SHA512_COMMIT_KEY}
+    # WHEN encrypting and then decrypting the encrypted data
+    encrypted_data = data_masker.encrypt(data, fields=fields_operation, provider_options=provider_options)
+    decrypted_data = data_masker.decrypt(encrypted_data, fields=fields_operation)
+
+    # THEN the result is only the specified fields are masked
+    assert decrypted_data == json.loads(data)