diff --git a/misk-aws2-dynamodb/api/misk-aws2-dynamodb.api b/misk-aws2-dynamodb/api/misk-aws2-dynamodb.api index 9a73bb5b33c..7ce55908c4a 100644 --- a/misk-aws2-dynamodb/api/misk-aws2-dynamodb.api +++ b/misk-aws2-dynamodb/api/misk-aws2-dynamodb.api @@ -71,6 +71,9 @@ public final class misk/aws2/dynamodb/testing/DynamoDbTable { public final class misk/aws2/dynamodb/testing/ExternalTestDynamoDbClientModule : misk/inject/KAbstractModule { public fun (ILjava/util/List;)V + public fun (ILjava/util/List;Lmisk/aws2/dynamodb/testing/ResetStrategy;)V + public synthetic fun (ILjava/util/List;Lmisk/aws2/dynamodb/testing/ResetStrategy;ILkotlin/jvm/internal/DefaultConstructorMarker;)V + public fun (ILmisk/aws2/dynamodb/testing/ResetStrategy;[Lmisk/aws2/dynamodb/testing/DynamoDbTable;)V public fun (I[Lmisk/aws2/dynamodb/testing/DynamoDbTable;)V public final fun provideRequiredTables ()Ljava/util/List; public final fun providesAmazonDynamoDB (Lapp/cash/tempest2/testing/TestDynamoDbClient;)Lsoftware/amazon/awssdk/services/dynamodb/DynamoDbClient; @@ -89,6 +92,14 @@ public final class misk/aws2/dynamodb/testing/ParallelTestsTableNameMapper : mis public fun mapName (Ljava/lang/String;)Ljava/lang/String; } +public final class misk/aws2/dynamodb/testing/ResetStrategy : java/lang/Enum { + public static final field DROP_RECREATE Lmisk/aws2/dynamodb/testing/ResetStrategy; + public static final field TRUNCATE Lmisk/aws2/dynamodb/testing/ResetStrategy; + public static fun getEntries ()Lkotlin/enums/EnumEntries; + public static fun valueOf (Ljava/lang/String;)Lmisk/aws2/dynamodb/testing/ResetStrategy; + public static fun values ()[Lmisk/aws2/dynamodb/testing/ResetStrategy; +} + public final class misk/aws2/dynamodb/testing/TestDynamoDb : com/google/common/util/concurrent/Service, misk/testing/TestFixture { public fun (Lapp/cash/tempest2/testing/internal/TestDynamoDbService;)V public fun addListener (Lcom/google/common/util/concurrent/Service$Listener;Ljava/util/concurrent/Executor;)V diff --git a/misk-aws2-dynamodb/src/test/kotlin/misk/aws2/dynamodb/testing/ExternalTestDynamoDbClientModuleTest.kt b/misk-aws2-dynamodb/src/test/kotlin/misk/aws2/dynamodb/testing/ExternalTestDynamoDbClientModuleTest.kt new file mode 100644 index 00000000000..4733e58f277 --- /dev/null +++ b/misk-aws2-dynamodb/src/test/kotlin/misk/aws2/dynamodb/testing/ExternalTestDynamoDbClientModuleTest.kt @@ -0,0 +1,151 @@ +package misk.aws2.dynamodb.testing + +import app.cash.tempest2.testing.JvmDynamoDbServer +import app.cash.tempest2.testing.TestDynamoDbClient +import jakarta.inject.Inject +import java.net.ServerSocket +import java.time.LocalDate +import misk.MiskTestingServiceModule +import misk.inject.KAbstractModule +import misk.testing.MiskTest +import misk.testing.MiskTestModule +import misk.testing.TestFixture +import org.assertj.core.api.Assertions.assertThat +import org.junit.jupiter.api.AfterAll +import org.junit.jupiter.api.Test +import org.junit.jupiter.api.TestInstance +import software.amazon.awssdk.enhanced.dynamodb.DynamoDbEnhancedClient +import software.amazon.awssdk.enhanced.dynamodb.model.EnhancedGlobalSecondaryIndex +import software.amazon.awssdk.enhanced.dynamodb.model.QueryConditional +import software.amazon.awssdk.services.dynamodb.model.ProjectionType + +@MiskTest(startService = true) +@TestInstance(TestInstance.Lifecycle.PER_CLASS) +class ExternalTestDynamoDbClientModuleTest { + private val server = + JvmDynamoDbServer.Factory.create(ServerSocket(0).use { it.localPort }).also { it.startAsync().awaitRunning() } + + @MiskTestModule val module = ExternalDynamoTestModule(server.port) + + @Inject lateinit var client: TestDynamoDbClient + @Inject lateinit var testFixtures: Set + + @AfterAll + fun stopServer() { + server.stopAsync().awaitTerminated() + } + + @Test + fun `reset truncates tables and indexes`() { + val movieTable = + DynamoDbEnhancedClient.builder() + .dynamoDbClient(client.dynamoDb) + .build() + .table(client.tables.single().tableName, AbstractDynamoDbTest.MOVIE_TABLE_SCHEMA) + val tableCreatedAt = + client.dynamoDb.describeTable { it.tableName(client.tables.single().tableName) }.table().creationDateTime() + + assertThat(movieTable.scan().items()).isEmpty() + + repeat(30) { index -> + movieTable.putItem( + DyMovie().apply { + name = "Movie $index" + release_date = LocalDate.of(2000, 1, 1).plusDays(index.toLong()) + directed_by = "Director" + } + ) + } + assertThat(movieTable.scan().items()).hasSize(30) + assertThat( + movieTable + .index("movies.release_date_index") + .query(QueryConditional.keyEqualTo { it.partitionValue("Director") }) + .flatMap { it.items() } + ) + .hasSize(30) + + testFixtures.single { it.javaClass.simpleName == "TestDynamoDbFixture" }.reset() + + assertThat( + client.dynamoDb.describeTable { it.tableName(client.tables.single().tableName) }.table().creationDateTime() + ) + .isEqualTo(tableCreatedAt) + assertThat(movieTable.scan().items()).isEmpty() + assertThat( + movieTable + .index("movies.release_date_index") + .query(QueryConditional.keyEqualTo { it.partitionValue("Director") }) + .flatMap { it.items() } + ) + .isEmpty() + } +} + +@MiskTest(startService = true) +@TestInstance(TestInstance.Lifecycle.PER_CLASS) +class DropRecreateExternalTestDynamoDbClientModuleTest { + private val server = + JvmDynamoDbServer.Factory.create(ServerSocket(0).use { it.localPort }).also { it.startAsync().awaitRunning() } + + @MiskTestModule val module = ExternalDynamoTestModule(server.port, ResetStrategy.DROP_RECREATE) + + @Inject lateinit var client: TestDynamoDbClient + @Inject lateinit var testFixtures: Set + + @AfterAll + fun stopServer() { + server.stopAsync().awaitTerminated() + } + + @Test + fun `reset drops and recreates tables`() { + val tableName = client.tables.single().tableName + val movieTable = + DynamoDbEnhancedClient.builder() + .dynamoDbClient(client.dynamoDb) + .build() + .table(tableName, AbstractDynamoDbTest.MOVIE_TABLE_SCHEMA) + val tableCreatedAt = client.dynamoDb.describeTable { it.tableName(tableName) }.table().creationDateTime() + + movieTable.putItem( + DyMovie().apply { + name = "Movie" + release_date = LocalDate.of(2000, 1, 1) + directed_by = "Director" + } + ) + + testFixtures.single { it.javaClass.simpleName == "TestDynamoDbFixture" }.reset() + + assertThat(client.dynamoDb.describeTable { it.tableName(tableName) }.table().creationDateTime()) + .isNotEqualTo(tableCreatedAt) + assertThat(movieTable.scan().items()).isEmpty() + } +} + +class ExternalDynamoTestModule( + private val port: Int, + private val resetStrategy: ResetStrategy = ResetStrategy.TRUNCATE, +) : KAbstractModule() { + override fun configure() { + install(MiskTestingServiceModule()) + install( + ExternalTestDynamoDbClientModule( + port, + listOf( + DynamoDbTable("movies", DyMovie::class) { createTableEnhancedRequest -> + createTableEnhancedRequest.globalSecondaryIndices( + EnhancedGlobalSecondaryIndex.builder() + .indexName("movies.release_date_index") + .projection { it.projectionType(ProjectionType.ALL) } + .provisionedThroughput { it.readCapacityUnits(40_000L).writeCapacityUnits(40_000L) } + .build() + ) + } + ), + resetStrategy, + ) + ) + } +} diff --git a/misk-aws2-dynamodb/src/testFixtures/kotlin/misk/aws2/dynamodb/testing/ExternalTestDynamoDbClientModule.kt b/misk-aws2-dynamodb/src/testFixtures/kotlin/misk/aws2/dynamodb/testing/ExternalTestDynamoDbClientModule.kt index 0910c553625..f6db3e145a3 100644 --- a/misk-aws2-dynamodb/src/testFixtures/kotlin/misk/aws2/dynamodb/testing/ExternalTestDynamoDbClientModule.kt +++ b/misk-aws2-dynamodb/src/testFixtures/kotlin/misk/aws2/dynamodb/testing/ExternalTestDynamoDbClientModule.kt @@ -9,6 +9,7 @@ import com.google.common.util.concurrent.AbstractService import com.google.inject.Provides import jakarta.inject.Inject import jakarta.inject.Singleton +import java.util.concurrent.ConcurrentHashMap import misk.ServiceModule import misk.aws2.dynamodb.DynamoDbService import misk.aws2.dynamodb.RequiredDynamoDbTable @@ -17,25 +18,51 @@ import misk.inject.KAbstractModule import misk.testing.TestFixture import misk.testing.updateForParallelTests import software.amazon.awssdk.services.dynamodb.DynamoDbClient +import software.amazon.awssdk.services.dynamodb.model.AttributeValue +import software.amazon.awssdk.services.dynamodb.model.DeleteRequest import software.amazon.awssdk.services.dynamodb.model.DeleteTableRequest import software.amazon.awssdk.services.dynamodb.model.ResourceNotFoundException +import software.amazon.awssdk.services.dynamodb.model.WriteRequest import software.amazon.awssdk.services.dynamodb.streams.DynamoDbStreamsClient +enum class ResetStrategy { + TRUNCATE, + DROP_RECREATE, +} + /** * Unlike the InProcessDynamoDbModule and DockerDynamoDbModule classes, this module does not internally start DynamoDB, * and instead relies on an external DynamoDB server already running on the given port, e.g. using a Gradle task as a * dependency of the test task, which starts DynamoDB in a Docker container. + * + * By default, tables are created once per injector and rows are truncated on later resets to avoid slow DynamoDB DDL + * calls before every test. Use [ResetStrategy.DROP_RECREATE] if tests use DynamoDB Streams or modify table definitions: + * truncation preserves streams (including delete records from the reset) and does not undo changes to TTL, billing + * mode, or indexes. */ -class ExternalTestDynamoDbClientModule(private val port: Int, originalTables: List) : KAbstractModule() { +class ExternalTestDynamoDbClientModule +@JvmOverloads +constructor( + private val port: Int, + originalTables: List, + private val resetStrategy: ResetStrategy = ResetStrategy.TRUNCATE, +) : KAbstractModule() { private val tables = originalTables.map { it.copy(tableName = ParallelTestsTableNameMapper.mapName(it.tableName)) } constructor(port: Int, vararg tables: DynamoDbTable) : this(port, tables.toList()) + constructor( + port: Int, + resetStrategy: ResetStrategy, + vararg tables: DynamoDbTable, + ) : this(port, tables.toList(), resetStrategy) + override fun configure() { for (table in tables) { multibind().toInstance(table) } + bind().toInstance(resetStrategy) bind().to() install(ServiceModule().dependsOn()) install(ServiceModule()) @@ -71,8 +98,14 @@ class ExternalTestDynamoDbClientModule(private val port: Int, originalTables: Li @Singleton private class TestDynamoDbFixture @Inject -constructor(private val client: TestDynamoDbClient, private val tables: List) : - AbstractIdleService(), TestFixture { +constructor( + private val client: TestDynamoDbClient, + private val tables: List, + private val resetStrategy: ResetStrategy, +) : AbstractIdleService(), TestFixture { + + private var tablesCreated = false + private val keyAttributeNames = ConcurrentHashMap>() override fun startUp() { reset() @@ -81,20 +114,77 @@ constructor(private val client: TestDynamoDbClient, private val tables: List - TestTable.create(table.tableName, table.tableClass) { table.configureTable(it.toBuilder()).build() } - }) { - client.dynamoDb.createTable(table) + + client.tables.forEach { client.dynamoDb.createTable(it) } + } + + private fun truncate(table: TestTable) { + val tableName = table.tableName + val keyNames = + keyAttributeNames.getOrPut(tableName) { + client.dynamoDb.describeTable { it.tableName(tableName) }.table().keySchema().map { it.attributeName() } + } + val keyAliases = keyNames.withIndex().associate { (index, name) -> "#k$index" to name } + var exclusiveStartKey: Map? = null + + do { + val startKey = exclusiveStartKey + val scan = + client.dynamoDb.scan { + it + .tableName(tableName) + .projectionExpression(keyAliases.keys.joinToString(", ")) + .expressionAttributeNames(keyAliases) + .consistentRead(true) + if (startKey != null) it.exclusiveStartKey(startKey) + } + deleteItems(tableName, scan.items()) + exclusiveStartKey = if (scan.hasLastEvaluatedKey()) scan.lastEvaluatedKey() else null + } while (exclusiveStartKey != null) + } + + private fun deleteItems(tableName: String, keys: List>) { + for (chunk in keys.chunked(MAX_BATCH_WRITE_ITEMS)) { + var writes = + chunk.map { key -> WriteRequest.builder().deleteRequest(DeleteRequest.builder().key(key).build()).build() } + var attempts = 0 + while (writes.isNotEmpty()) { + if (attempts > 0) Thread.sleep(UNPROCESSED_RETRY_DELAY_MS) + check(attempts < MAX_BATCH_WRITE_ATTEMPTS) { + "BatchWriteItem left ${writes.size} unprocessed delete(s) for table $tableName after " + + "$MAX_BATCH_WRITE_ATTEMPTS attempts" + } + val response = client.dynamoDb.batchWriteItem { it.requestItems(mapOf(tableName to writes)) } + writes = response.unprocessedItems()[tableName].orEmpty() + attempts++ + } } } + + private companion object { + const val MAX_BATCH_WRITE_ITEMS = 25 + const val MAX_BATCH_WRITE_ATTEMPTS = 5 + const val UNPROCESSED_RETRY_DELAY_MS = 50L + } } @Singleton