auto commit

Author: rishav-karanjit

Examples/runtimes/java/Migration/PlaintextToAWSDBE/src/test/java/software/amazon/cryptography/examples/migration/TestUtils.java

@@ -1,5 +1,9 @@
 package software.amazon.cryptography.examples.migration;
 
+import java.util.HashMap;
+import software.amazon.awssdk.services.dynamodb.DynamoDbClient;
+import software.amazon.awssdk.services.dynamodb.model.AttributeValue;
+import software.amazon.awssdk.services.dynamodb.model.DeleteItemRequest;
 import java.util.UUID;
 
 public class TestUtils {
@@ -14,4 +18,40 @@ public class TestUtils {
 
   public static final String PARTITION_KEY =
     "PlaintextMigrationExample" + UUID.randomUUID();
+  
+  /**
+   * Deletes an item from a DynamoDB table.
+   *
+   * @param tableName The name of the DynamoDB table
+   * @param partitionKeyName The name of partition key
+   * @param sortKeyName The name of sort key
+   * @param partitionKeyValue The value of the partition key
+   * @param sortKeyValue The value of the sort key (can be null if table doesn't have a sort key)
+   */
+  public static void cleanUpDDBItem(
+    final String tableName,
+    final String partitionKeyName,
+    final String sortKeyName,
+    final String partitionKeyValue,
+    final String sortKeyValue
+  ) {
+    final DynamoDbClient ddb = DynamoDbClient.builder().build();
+    final HashMap<String, AttributeValue> keyToDelete = new HashMap<>();
+    keyToDelete.put(
+      partitionKeyName,
+      AttributeValue.builder().s(partitionKeyValue).build()
+    );
+    if (sortKeyValue != null) {
+      keyToDelete.put(
+        sortKeyName,
+        AttributeValue.builder().n(sortKeyValue).build()
+      );
+    }
+    final DeleteItemRequest deleteRequest = DeleteItemRequest
+      .builder()
+      .tableName(tableName)
+      .key(keyToDelete)
+      .build();
+    ddb.deleteItem(deleteRequest);
+  }
 }

Examples/runtimes/java/Migration/PlaintextToAWSDBE/src/test/java/software/amazon/cryptography/examples/migration/awsdbe/TestMigrationExampleStep1.java

@@ -1,5 +1,7 @@
 package software.amazon.cryptography.examples.migration.awsdbe;
 
+import java.util.Arrays;
+import java.util.List;
 import org.testng.annotations.Test;
 import software.amazon.cryptography.examples.awsdbe.MigrationExampleStep1;
 import software.amazon.cryptography.examples.awsdbe.MigrationExampleStep2;
@@ -62,5 +64,9 @@ public void TestMigrationStep1() {
       3,
       TestUtils.PARTITION_KEY
     );
+    List<String> sortkeys = Arrays.asList("0", "1", "2", "3");
+    for (String sortkey : sortkeys) {
+      TestUtils.cleanUpDDBItem(TestUtils.TEST_DDB_TABLE_NAME, "partition_key", "sort_key", TestUtils.PARTITION_KEY, sortkey);
+    }
   }
 }

Examples/runtimes/java/Migration/PlaintextToAWSDBE/src/test/java/software/amazon/cryptography/examples/migration/awsdbe/TestMigrationExampleStep2.java

@@ -1,5 +1,7 @@
 package software.amazon.cryptography.examples.migration.awsdbe;
 
+import java.util.Arrays;
+import java.util.List;
 import org.testng.annotations.Test;
 import software.amazon.cryptography.examples.awsdbe.MigrationExampleStep1;
 import software.amazon.cryptography.examples.awsdbe.MigrationExampleStep2;
@@ -62,5 +64,9 @@ public void TestMigrationStep2() {
       3,
       TestUtils.PARTITION_KEY
     );
+    List<String> sortkeys = Arrays.asList("0", "1", "2", "3");
+    for (String sortkey : sortkeys) {
+      TestUtils.cleanUpDDBItem(TestUtils.TEST_DDB_TABLE_NAME, "partition_key", "sort_key", TestUtils.PARTITION_KEY, sortkey);
+    }
   }
 }

Examples/runtimes/java/Migration/PlaintextToAWSDBE/src/test/java/software/amazon/cryptography/examples/migration/awsdbe/TestMigrationExampleStep3.java

@@ -1,7 +1,8 @@
 package software.amazon.cryptography.examples.migration.awsdbe;
 
 import static org.testng.Assert.assertThrows;
-
+import java.util.Arrays;
+import java.util.List;
 import org.testng.annotations.Test;
 import software.amazon.awssdk.core.exception.SdkClientException;
 import software.amazon.cryptography.examples.awsdbe.MigrationExampleStep1;
@@ -75,5 +76,9 @@ public void TestMigrationStep0() {
       2,
       TestUtils.PARTITION_KEY
     );
+    List<String> sortkeys = Arrays.asList("0", "1", "2", "3");
+    for (String sortkey : sortkeys) {
+      TestUtils.cleanUpDDBItem(TestUtils.TEST_DDB_TABLE_NAME, "partition_key", "sort_key", TestUtils.PARTITION_KEY, sortkey);
+    }
   }
 }

Examples/runtimes/java/Migration/PlaintextToAWSDBE/src/test/java/software/amazon/cryptography/examples/migration/awsdbe/encrypttable/TestEncryptExistingTable.java

@@ -204,5 +204,10 @@ public void TestEncryptExistingTable() {
       TestUtils.TEST_KMS_KEY_ID,
       TestUtils.TEST_DDB_TABLE_NAME
     );
+
+    List<String> sortkeys = Arrays.asList("0", "1", "2");
+    for (String sortkey : sortkeys) {
+      TestUtils.cleanUpDDBItem(TestUtils.TEST_DDB_TABLE_NAME, "partition_key", "sort_key", TestUtils.PARTITION_KEY, sortkey);
+    }
   }
 }

Examples/runtimes/java/Migration/PlaintextToAWSDBE/src/test/java/software/amazon/cryptography/examples/migration/plaintext/TestMigrationExampleStep0.java

@@ -1,7 +1,8 @@
 package software.amazon.cryptography.examples.migration.plaintext;
 
 import static org.testng.Assert.assertThrows;
-
+import java.util.Arrays;
+import java.util.List;
 import org.testng.annotations.Test;
 import software.amazon.cryptography.examples.awsdbe.MigrationExampleStep1;
 import software.amazon.cryptography.examples.awsdbe.MigrationExampleStep2;
@@ -71,5 +72,9 @@ public void TestMigrationStep0() {
         );
       }
     );
+    List<String> sortkeys = Arrays.asList("0", "1", "2", "3");
+    for (String sortkey : sortkeys) {
+      TestUtils.cleanUpDDBItem(TestUtils.TEST_DDB_TABLE_NAME, "partition_key", "sort_key", TestUtils.PARTITION_KEY, sortkey);
+    }
   }
 }