From fa1e151fed1ca3dcd1fff14fe13af1c95d3fe61d Mon Sep 17 00:00:00 2001
From: Rishav karanjit <karanjitrishav4@gmail.com>
Date: Tue, 7 Apr 2026 12:10:04 -0700
Subject: [PATCH] chore(java): shut down local DDB in test (#2176)

---
 .../datamodeling/TransformerHolisticIT.java   |  9 ++++
 .../CachingMostRecentProviderTests.java       | 43 ++++++++++++++-----
 .../mapper/integration/DynamoDBTestBase.java  |  9 ++++
 3 files changed, 50 insertions(+), 11 deletions(-)

diff --git a/DynamoDbEncryption/runtimes/java/src/test/sdkv1/com/amazonaws/services/dynamodbv2/datamodeling/TransformerHolisticIT.java b/DynamoDbEncryption/runtimes/java/src/test/sdkv1/com/amazonaws/services/dynamodbv2/datamodeling/TransformerHolisticIT.java
index c37fdf9643..0ece86494f 100644
--- a/DynamoDbEncryption/runtimes/java/src/test/sdkv1/com/amazonaws/services/dynamodbv2/datamodeling/TransformerHolisticIT.java
+++ b/DynamoDbEncryption/runtimes/java/src/test/sdkv1/com/amazonaws/services/dynamodbv2/datamodeling/TransformerHolisticIT.java
@@ -55,6 +55,7 @@
 import java.util.stream.Collectors;
 import javax.crypto.SecretKey;
 import javax.crypto.spec.SecretKeySpec;
+import org.testng.annotations.AfterTest;
 import org.testng.annotations.BeforeTest;
 import org.testng.annotations.DataProvider;
 import org.testng.annotations.Test;
@@ -326,6 +327,14 @@ public void setUp() {
     createCiphertextTables(client);
   }
 
+  @AfterTest
+  public void tearDownDDBLocal() {
+    if (client != null) {
+      client.shutdown();
+      client = null;
+    }
+  }
+
   @Test(dataProvider = "getDecryptTestVectors")
   public void decryptTestVector(Scenario scenario) throws IOException {
     client = DynamoDBEmbedded.create().amazonDynamoDB();
diff --git a/DynamoDbEncryption/runtimes/java/src/test/sdkv1/com/amazonaws/services/dynamodbv2/datamodeling/encryption/providers/CachingMostRecentProviderTests.java b/DynamoDbEncryption/runtimes/java/src/test/sdkv1/com/amazonaws/services/dynamodbv2/datamodeling/encryption/providers/CachingMostRecentProviderTests.java
index b0b8c24839..ee82d1b012 100644
--- a/DynamoDbEncryption/runtimes/java/src/test/sdkv1/com/amazonaws/services/dynamodbv2/datamodeling/encryption/providers/CachingMostRecentProviderTests.java
+++ b/DynamoDbEncryption/runtimes/java/src/test/sdkv1/com/amazonaws/services/dynamodbv2/datamodeling/encryption/providers/CachingMostRecentProviderTests.java
@@ -26,6 +26,9 @@
 import java.util.Map;
 import javax.crypto.SecretKey;
 import javax.crypto.spec.SecretKeySpec;
+import org.testng.annotations.AfterClass;
+import org.testng.annotations.AfterMethod;
+import org.testng.annotations.BeforeClass;
 import org.testng.annotations.BeforeMethod;
 import org.testng.annotations.Test;
 
@@ -47,30 +50,48 @@ public class CachingMostRecentProviderTests {
   private static final DynamoDBEncryptor ENCRYPTOR =
     DynamoDBEncryptor.getInstance(BASE_PROVIDER);
 
-  private AmazonDynamoDB client;
-  private Map<String, Integer> methodCalls;
+  private static AmazonDynamoDB rawClient;
+  private static AmazonDynamoDB instrumentedClient;
+  private static Map<String, Integer> methodCalls;
   private ProviderStore store;
   private EncryptionContext ctx;
 
+  // DDB Local client is created once per class since DynamoDBEmbedded.create() is expensive.
+  // Table is created/deleted per method to ensure test isolation, as tests assert
+  // exact DynamoDB call counts and expect a clean table with no leftover data.
+  @BeforeClass
+  public static void setupClient() {
+    methodCalls = new HashMap<String, Integer>();
+    rawClient = DynamoDBEmbedded.create().amazonDynamoDB();
+    instrumentedClient =
+      instrument(rawClient, AmazonDynamoDB.class, methodCalls);
+  }
+
+  @AfterClass
+  public static void tearDownDDBLocal() {
+    if (rawClient != null) {
+      rawClient.shutdown();
+      rawClient = null;
+    }
+  }
+
   @BeforeMethod
   public void setup() {
-    methodCalls = new HashMap<String, Integer>();
-    client =
-      instrument(
-        DynamoDBEmbedded.create().amazonDynamoDB(),
-        AmazonDynamoDB.class,
-        methodCalls
-      );
     MetaStore.createTable(
-      client,
+      instrumentedClient,
       TABLE_NAME,
       new ProvisionedThroughput(1L, 1L)
     );
-    store = new MetaStore(client, TABLE_NAME, ENCRYPTOR);
+    store = new MetaStore(instrumentedClient, TABLE_NAME, ENCRYPTOR);
     ctx = new EncryptionContext.Builder().build();
     methodCalls.clear();
   }
 
+  @AfterMethod
+  public void tearDown() {
+    rawClient.deleteTable(TABLE_NAME);
+  }
+
   @Test
   public void testConstructors() {
     final CachingMostRecentProvider prov = new CachingMostRecentProvider(
diff --git a/DynamoDbEncryption/runtimes/java/src/test/sdkv1/com/amazonaws/services/dynamodbv2/mapper/integration/DynamoDBTestBase.java b/DynamoDbEncryption/runtimes/java/src/test/sdkv1/com/amazonaws/services/dynamodbv2/mapper/integration/DynamoDBTestBase.java
index 6c4be28481..ad2b485e2f 100644
--- a/DynamoDbEncryption/runtimes/java/src/test/sdkv1/com/amazonaws/services/dynamodbv2/mapper/integration/DynamoDBTestBase.java
+++ b/DynamoDbEncryption/runtimes/java/src/test/sdkv1/com/amazonaws/services/dynamodbv2/mapper/integration/DynamoDBTestBase.java
@@ -21,6 +21,7 @@
 import java.util.Collection;
 import java.util.HashSet;
 import java.util.Set;
+import org.testng.annotations.AfterClass;
 
 public class DynamoDBTestBase {
 
@@ -30,6 +31,14 @@ public static void setUpTestBase() {
     dynamo = DynamoDBEmbedded.create().amazonDynamoDB();
   }
 
+  @AfterClass
+  public static void tearDownDDBLocal() {
+    if (dynamo != null) {
+      dynamo.shutdown();
+      dynamo = null;
+    }
+  }
+
   public static AmazonDynamoDB getClient() {
     if (dynamo == null) {
       setUpTestBase();