Add Cleanup Utility

Author: imabhichow

Examples/runtimes/python/DynamoDBEncryption/test/cleanup.py

@@ -0,0 +1,76 @@
+# Copyright Amazon.com Inc. or its affiliates. All Rights Reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""
+Test cleanup utilities for DynamoDB Encryption SDK.
+
+This module provides utilities for cleaning up resources after running tests.
+NOTE: This is only a test utility and should not be used in production code.
+It is specifically designed for cleaning up test resources after test execution.
+"""
+import boto3
+
+BRANCH_KEY_IDENTIFIER_FIELD = "branch-key-id"
+TYPE_FIELD = "type"
+
+
+def delete_branch_key(
+    identifier: str,
+    table_name: str,
+    ddb_client: boto3.client,
+) -> bool:
+    """
+    Delete all branch key items with the given identifier.
+
+    Args:
+        identifier: Branch key identifier to delete
+        table_name: DynamoDB table name
+        ddb_client: DynamoDB client to use
+
+    Returns:
+        True if all items were deleted, False if more than 100 items exist
+
+    Raises:
+        ValueError: If an item is not a branch key
+
+    """
+    if ddb_client is None:
+        ddb_client = boto3.client("dynamodb")
+
+    # Query for items with matching identifier
+    query_response = ddb_client.query(
+        TableName=table_name,
+        KeyConditionExpression="#pk = :pk",
+        ExpressionAttributeNames={"#pk": BRANCH_KEY_IDENTIFIER_FIELD},
+        ExpressionAttributeValues={":pk": {"S": identifier}},
+    )
+
+    items = query_response.get("Items", [])
+    if not items:
+        return True
+
+    # Create delete requests for each item
+    delete_items = []
+    for item in items:
+        if TYPE_FIELD not in item:
+            raise ValueError("Item is not a branch key")
+
+        delete_item = {
+            "Delete": {
+                "Key": {BRANCH_KEY_IDENTIFIER_FIELD: {"S": identifier}, TYPE_FIELD: item[TYPE_FIELD]},
+                "TableName": table_name,
+            }
+        }
+        delete_items.append(delete_item)
+
+    if not delete_items:
+        return True
+
+    # DynamoDB transactions are limited to 100 items
+    if len(delete_items) > 100:
+        delete_items = delete_items[:100]
+
+    # Execute the delete transaction
+    ddb_client.transact_write_items(TransactItems=delete_items)
+
+    # Return False if we had to truncate the deletion
+    return len(items) <= 100

Examples/runtimes/python/DynamoDBEncryption/test/keyring/test_hierarchical_keyring_example.py

@@ -7,6 +7,7 @@
 
 from ...src.create_keystore_key_example import keystore_create_key
 from ...src.keyring.hierarchical_keyring_example import hierarchical_keyring_get_item_put_item
+from ..cleanup import delete_branch_key
 from ..test_utils import (
     TEST_DDB_TABLE_NAME,
     TEST_KEYSTORE_KMS_KEY_ID,
@@ -30,3 +31,7 @@ def test_hierarchical_keyring_example():
     hierarchical_keyring_get_item_put_item(
         TEST_DDB_TABLE_NAME, key_id1, key_id2, TEST_KEYSTORE_NAME, TEST_LOGICAL_KEYSTORE_NAME, TEST_KEYSTORE_KMS_KEY_ID
     )
+
+    # Cleanup Branch Key
+    delete_branch_key(key_id1, TEST_KEYSTORE_NAME, None)
+    delete_branch_key(key_id2, TEST_KEYSTORE_NAME, None)

Examples/runtimes/python/DynamoDBEncryption/test/keyring/test_shared_cache_across_hierarchical_keyrings_example.py

@@ -9,6 +9,7 @@
 from ...src.keyring.shared_cache_across_hierarchical_keyrings_example import (
     shared_cache_across_hierarchical_keyrings_example,
 )
+from ..cleanup import delete_branch_key
 from ..test_utils import (
     TEST_DDB_TABLE_NAME,
     TEST_KEYSTORE_KMS_KEY_ID,
@@ -37,3 +38,6 @@ def test_shared_cache_across_hierarchical_keyrings_example():
         TEST_PARTITION_ID,
         TEST_KEYSTORE_KMS_KEY_ID,
     )
+
+    # Cleanup Branch Key
+    delete_branch_key(key_id, TEST_KEYSTORE_NAME, None)

Examples/runtimes/python/DynamoDBEncryption/test/test_create_keystore_key_example.py

@@ -4,6 +4,7 @@
 import pytest
 
 from ..src.create_keystore_key_example import keystore_create_key
+from .cleanup import delete_branch_key
 from .test_utils import TEST_KEYSTORE_KMS_KEY_ID, TEST_KEYSTORE_NAME, TEST_LOGICAL_KEYSTORE_NAME
 
 pytestmark = [pytest.mark.examples]
@@ -14,4 +15,6 @@ def test_create_keystore_key_example():
     key_id = keystore_create_key(TEST_KEYSTORE_NAME, TEST_LOGICAL_KEYSTORE_NAME, TEST_KEYSTORE_KMS_KEY_ID)
 
     assert key_id is not None
-    # TODO: Cleanup Branch Key Items
+
+    # Cleanup Branch Key
+    delete_branch_key(key_id, TEST_KEYSTORE_NAME, None)