From 64f40eff06ba7bbd3dc7369c4d0122993d5e192d Mon Sep 17 00:00:00 2001
From: Mehdi Mollaverdi <mmollaverdi@squareup.com>
Date: Wed, 7 Aug 2024 10:23:16 +1000
Subject: [PATCH] Add method to test DDB clients for resetting tables in
 between tests (#189)

---
 .../testing/internal/DefaultTestDynamoDbClient.kt  | 14 +++++++++-----
 .../app/cash/tempest/testing/TestDynamoDbClient.kt |  3 +++
 .../testing/internal/DefaultTestDynamoDbClient.kt  | 14 +++++++++-----
 .../cash/tempest2/testing/TestDynamoDbClient.kt    |  3 +++
 4 files changed, 24 insertions(+), 10 deletions(-)

diff --git a/tempest-testing-internal/src/main/kotlin/app/cash/tempest/testing/internal/DefaultTestDynamoDbClient.kt b/tempest-testing-internal/src/main/kotlin/app/cash/tempest/testing/internal/DefaultTestDynamoDbClient.kt
index 73da471cc..2a52f8ed6 100644
--- a/tempest-testing-internal/src/main/kotlin/app/cash/tempest/testing/internal/DefaultTestDynamoDbClient.kt
+++ b/tempest-testing-internal/src/main/kotlin/app/cash/tempest/testing/internal/DefaultTestDynamoDbClient.kt
@@ -31,6 +31,15 @@ class DefaultTestDynamoDbClient(
   override val dynamoDbStreams = buildDynamoDbStreams(hostName, port)
 
   override fun startUp() {
+    reset()
+  }
+
+  override fun shutDown() {
+    dynamoDb.shutdown()
+    dynamoDbStreams.shutdown()
+  }
+
+  override fun reset() {
     // Cleans up the tables before each run.
     for (tableName in dynamoDb.listTables().tableNames) {
       dynamoDb.deleteTable(tableName)
@@ -39,9 +48,4 @@ class DefaultTestDynamoDbClient(
       dynamoDb.createTable(table)
     }
   }
-
-  override fun shutDown() {
-    dynamoDb.shutdown()
-    dynamoDbStreams.shutdown()
-  }
 }
diff --git a/tempest-testing/src/main/kotlin/app/cash/tempest/testing/TestDynamoDbClient.kt b/tempest-testing/src/main/kotlin/app/cash/tempest/testing/TestDynamoDbClient.kt
index b0726a21d..9eb84d2a2 100644
--- a/tempest-testing/src/main/kotlin/app/cash/tempest/testing/TestDynamoDbClient.kt
+++ b/tempest-testing/src/main/kotlin/app/cash/tempest/testing/TestDynamoDbClient.kt
@@ -48,6 +48,9 @@ interface TestDynamoDbClient : Service {
   fun <DB : LogicalDb> logicalDb(type: Class<DB>, mapperConfig: DynamoDBMapperConfig): DB {
     return logicalDb(type.kotlin, mapperConfig)
   }
+
+  /** Cleans up tables in between test runs. */
+  fun reset()
 }
 
 inline fun <reified DB : LogicalDb> TestDynamoDbClient.logicalDb(): DB {
diff --git a/tempest2-testing-internal/src/main/kotlin/app/cash/tempest2/testing/internal/DefaultTestDynamoDbClient.kt b/tempest2-testing-internal/src/main/kotlin/app/cash/tempest2/testing/internal/DefaultTestDynamoDbClient.kt
index 5f1b87266..a7a0b1c4f 100644
--- a/tempest2-testing-internal/src/main/kotlin/app/cash/tempest2/testing/internal/DefaultTestDynamoDbClient.kt
+++ b/tempest2-testing-internal/src/main/kotlin/app/cash/tempest2/testing/internal/DefaultTestDynamoDbClient.kt
@@ -34,6 +34,15 @@ class DefaultTestDynamoDbClient(
   override val asyncDynamoDbStreams = buildAsyncDynamoDbStreams(hostName, port)
 
   override fun startUp() {
+    reset()
+  }
+
+  override fun shutDown() {
+    dynamoDb.close()
+    dynamoDbStreams.close()
+  }
+
+  override fun reset() {
     // Cleans up the tables before each run.
     for (tableName in dynamoDb.listTables().tableNames()) {
       dynamoDb.deleteTable(DeleteTableRequest.builder().tableName(tableName).build())
@@ -42,9 +51,4 @@ class DefaultTestDynamoDbClient(
       dynamoDb.createTable(table)
     }
   }
-
-  override fun shutDown() {
-    dynamoDb.close()
-    dynamoDbStreams.close()
-  }
 }
diff --git a/tempest2-testing/src/main/kotlin/app/cash/tempest2/testing/TestDynamoDbClient.kt b/tempest2-testing/src/main/kotlin/app/cash/tempest2/testing/TestDynamoDbClient.kt
index caeb12b27..64f26e98d 100644
--- a/tempest2-testing/src/main/kotlin/app/cash/tempest2/testing/TestDynamoDbClient.kt
+++ b/tempest2-testing/src/main/kotlin/app/cash/tempest2/testing/TestDynamoDbClient.kt
@@ -128,6 +128,9 @@ interface TestDynamoDbClient : Service {
   ): DB {
     return asyncLogicalDb(type.kotlin, extensions)
   }
+
+  /** Cleans up tables in between test runs. */
+  fun reset()
 }
 
 inline fun <reified DB : LogicalDb> TestDynamoDbClient.logicalDb(