Skip to content

Commit b2a4d11

Browse files
mmollaverdisvc-squareup-copybara
mmollaverdi
authored and
svc-squareup-copybara
committed
Add support for running tests in parallel when dynamodb is
involved GitOrigin-RevId: f2a33bf4f3efd9741caec87886c1e40b81136093
1 parent 4bf5568 commit b2a4d11

File tree

7 files changed

+297
-1
lines changed

7 files changed

+297
-1
lines changed

‎misk-aws-dynamodb/api/misk-aws-dynamodb.api‎

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -49,6 +49,15 @@ public final class misk/aws/dynamodb/testing/TestDynamoDb : com/google/common/ut
4949
public fun stopAsync ()Lcom/google/common/util/concurrent/Service;
5050
}
5151

52+
public final class misk/aws/dynamodb/testing/TestDynamoDbClientModule : misk/inject/KAbstractModule {
53+
public fun <init> (ILjava/util/List;)V
54+
public fun <init> (I[Lmisk/aws/dynamodb/testing/DynamoDbTable;)V
55+
public final fun provideRequiredTables ()Ljava/util/List;
56+
public final fun providesAmazonDynamoDB (Lapp/cash/tempest/testing/TestDynamoDbClient;)Lcom/amazonaws/services/dynamodbv2/AmazonDynamoDB;
57+
public final fun providesAmazonDynamoDBStreams (Lapp/cash/tempest/testing/TestDynamoDbClient;)Lcom/amazonaws/services/dynamodbv2/AmazonDynamoDBStreams;
58+
public final fun providesTestDynamoDbClient ()Lapp/cash/tempest/testing/TestDynamoDbClient;
59+
}
60+
5261
public abstract interface class misk/dynamodb/DyTimestampedEntity {
5362
public abstract fun getCreated_at ()Ljava/util/Date;
5463
public abstract fun getUpdated_at ()Ljava/util/Date;

‎misk-aws-dynamodb/src/testFixtures/kotlin/misk/aws/dynamodb/testing/DynamoDbTable.kt‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -14,7 +14,7 @@ import kotlin.reflect.full.findAnnotation
1414
*/
1515
data class DynamoDbTable(
1616
val tableClass: KClass<*>,
17-
val configureTable: (CreateTableRequest) -> CreateTableRequest = { it }
17+
val configureTable: (CreateTableRequest) -> CreateTableRequest = { it },
1818
) {
1919
val tableName: String
2020
get() {

‎misk-aws-dynamodb/src/testFixtures/kotlin/misk/aws/dynamodb/testing/TestDynamoDbClientModule.kt‎

Lines changed: 63 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,63 @@
1+
package misk.aws.dynamodb.testing
2+
3+
import app.cash.tempest.testing.TestDynamoDbClient
4+
import app.cash.tempest.testing.TestTable
5+
import app.cash.tempest.testing.internal.DefaultTestDynamoDbClient
6+
import com.amazonaws.services.dynamodbv2.AmazonDynamoDB
7+
import com.amazonaws.services.dynamodbv2.AmazonDynamoDBStreams
8+
import com.google.common.util.concurrent.AbstractService
9+
import com.google.inject.Provides
10+
import jakarta.inject.Inject
11+
import jakarta.inject.Singleton
12+
import misk.ServiceModule
13+
import misk.dynamodb.DynamoDbService
14+
import misk.dynamodb.RequiredDynamoDbTable
15+
import misk.inject.KAbstractModule
16+
17+
class TestDynamoDbClientModule(
18+
private val port: Int,
19+
private val tables: List<DynamoDbTable>
20+
) : KAbstractModule() {
21+
22+
constructor(port: Int, vararg tables: DynamoDbTable) : this(port, tables.toList())
23+
24+
override fun configure() {
25+
for (table in tables) {
26+
multibind<DynamoDbTable>().toInstance(table)
27+
}
28+
bind<DynamoDbService>().to<TestDynamoDbService>()
29+
install(ServiceModule<DynamoDbService>())
30+
}
31+
32+
@Provides
33+
@Singleton
34+
fun provideRequiredTables(): List<RequiredDynamoDbTable> =
35+
tables.map { RequiredDynamoDbTable(it.tableName) }
36+
37+
@Provides
38+
@Singleton
39+
fun providesTestDynamoDbClient(): TestDynamoDbClient {
40+
return DefaultTestDynamoDbClient(
41+
tables = tables.map { TestTable.create(it.tableClass, it.configureTable) },
42+
port = port
43+
)
44+
}
45+
46+
@Provides
47+
@Singleton
48+
fun providesAmazonDynamoDB(testDynamoDbClient: TestDynamoDbClient): AmazonDynamoDB {
49+
return testDynamoDbClient.dynamoDb
50+
}
51+
52+
@Provides
53+
@Singleton
54+
fun providesAmazonDynamoDBStreams(testDynamoDbClient: TestDynamoDbClient): AmazonDynamoDBStreams {
55+
return testDynamoDbClient.dynamoDbStreams
56+
}
57+
}
58+
59+
@Singleton
60+
private class TestDynamoDbService @Inject constructor() : AbstractService(), DynamoDbService {
61+
override fun doStart() = notifyStarted()
62+
override fun doStop() = notifyStopped()
63+
}

‎misk-aws2-dynamodb/api/misk-aws2-dynamodb.api‎

Lines changed: 26 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -30,6 +30,18 @@ public final class misk/aws2/dynamodb/RequiredDynamoDbTable {
3030
public fun toString ()Ljava/lang/String;
3131
}
3232

33+
public abstract interface class misk/aws2/dynamodb/TableNameMapper {
34+
public abstract fun mapName (Ljava/lang/String;)Ljava/lang/String;
35+
}
36+
37+
public final class misk/aws2/dynamodb/TableNameMapper$DefaultImpls {
38+
public static fun mapName (Lmisk/aws2/dynamodb/TableNameMapper;Ljava/lang/String;)Ljava/lang/String;
39+
}
40+
41+
public final class misk/aws2/dynamodb/TableNameMapperKt {
42+
public static final fun withTableNameMapper (Lsoftware/amazon/awssdk/enhanced/dynamodb/DynamoDbEnhancedClient;Lmisk/aws2/dynamodb/TableNameMapper;)Lsoftware/amazon/awssdk/enhanced/dynamodb/DynamoDbEnhancedClient;
43+
}
44+
3345
public final class misk/aws2/dynamodb/testing/DockerDynamoDbModule : misk/inject/KAbstractModule {
3446
public fun <init> (Ljava/util/List;)V
3547
public fun <init> ([Lmisk/aws2/dynamodb/testing/DynamoDbTable;)V
@@ -56,6 +68,15 @@ public final class misk/aws2/dynamodb/testing/DynamoDbTable {
5668
public fun toString ()Ljava/lang/String;
5769
}
5870

71+
public final class misk/aws2/dynamodb/testing/ExternalTestDynamoDbClientModule : misk/inject/KAbstractModule {
72+
public fun <init> (ILjava/util/List;)V
73+
public fun <init> (I[Lmisk/aws2/dynamodb/testing/DynamoDbTable;)V
74+
public final fun provideRequiredTables ()Ljava/util/List;
75+
public final fun providesAmazonDynamoDB (Lapp/cash/tempest2/testing/TestDynamoDbClient;)Lsoftware/amazon/awssdk/services/dynamodb/DynamoDbClient;
76+
public final fun providesAmazonDynamoDBStreams (Lapp/cash/tempest2/testing/TestDynamoDbClient;)Lsoftware/amazon/awssdk/services/dynamodb/streams/DynamoDbStreamsClient;
77+
public final fun providesTestDynamoDbClient ()Lapp/cash/tempest2/testing/TestDynamoDbClient;
78+
}
79+
5980
public final class misk/aws2/dynamodb/testing/InProcessDynamoDbModule : misk/inject/KAbstractModule {
6081
public fun <init> (Ljava/util/List;)V
6182
public fun <init> ([Lmisk/aws2/dynamodb/testing/DynamoDbTable;)V
@@ -65,6 +86,11 @@ public final class misk/aws2/dynamodb/testing/InProcessDynamoDbModule : misk/inj
6586
public final fun providesTestDynamoDb ()Lmisk/aws2/dynamodb/testing/TestDynamoDb;
6687
}
6788

89+
public final class misk/aws2/dynamodb/testing/ParallelTestsTableNameMapper : misk/aws2/dynamodb/TableNameMapper {
90+
public static final field INSTANCE Lmisk/aws2/dynamodb/testing/ParallelTestsTableNameMapper;
91+
public fun mapName (Ljava/lang/String;)Ljava/lang/String;
92+
}
93+
6894
public final class misk/aws2/dynamodb/testing/TestDynamoDb : com/google/common/util/concurrent/Service, misk/testing/TestFixture {
6995
public fun <init> (Lapp/cash/tempest2/testing/internal/TestDynamoDbService;)V
7096
public fun addListener (Lcom/google/common/util/concurrent/Service$Listener;Ljava/util/concurrent/Executor;)V

‎misk-aws2-dynamodb/build.gradle.kts‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,7 @@ plugins {
99

1010
dependencies {
1111
api(libs.aws2Dynamodb)
12+
api(libs.aws2DynamodbEnhanced)
1213
api(libs.aws2Auth)
1314
api(libs.awsSdkCore)
1415
api(libs.guava)

‎misk-aws2-dynamodb/src/main/kotlin/misk/aws2/dynamodb/TableNameMapper.kt‎

Lines changed: 80 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,80 @@
1+
package misk.aws2.dynamodb
2+
3+
import software.amazon.awssdk.enhanced.dynamodb.Document
4+
import software.amazon.awssdk.enhanced.dynamodb.DynamoDbEnhancedClient
5+
import software.amazon.awssdk.enhanced.dynamodb.DynamoDbTable
6+
import software.amazon.awssdk.enhanced.dynamodb.TableSchema
7+
import software.amazon.awssdk.enhanced.dynamodb.model.BatchGetItemEnhancedRequest
8+
import software.amazon.awssdk.enhanced.dynamodb.model.BatchGetResultPageIterable
9+
import software.amazon.awssdk.enhanced.dynamodb.model.BatchWriteItemEnhancedRequest
10+
import software.amazon.awssdk.enhanced.dynamodb.model.BatchWriteResult
11+
import software.amazon.awssdk.enhanced.dynamodb.model.TransactGetItemsEnhancedRequest
12+
import software.amazon.awssdk.enhanced.dynamodb.model.TransactWriteItemsEnhancedRequest
13+
import software.amazon.awssdk.enhanced.dynamodb.model.TransactWriteItemsEnhancedResponse
14+
import java.util.function.Consumer
15+
16+
/**
17+
* Mapping dynamodb table names at runtime.
18+
* Used in parallel tests for using isolated tables per parallel test process.
19+
*/
20+
interface TableNameMapper {
21+
fun mapName(tableName: String): String = tableName
22+
}
23+
24+
fun DynamoDbEnhancedClient.withTableNameMapper(
25+
tableNameMapper: TableNameMapper,
26+
): DynamoDbEnhancedClient = MappedDynamoDbEnhancedClient(
27+
client = this,
28+
tableNameMapper = tableNameMapper,
29+
)
30+
31+
internal class MappedDynamoDbEnhancedClient(
32+
private val client: DynamoDbEnhancedClient,
33+
private val tableNameMapper: TableNameMapper,
34+
) : DynamoDbEnhancedClient {
35+
36+
override fun <T : Any?> table(name: String?, schema: TableSchema<T>?): DynamoDbTable<T> = client.table(
37+
name?.let { tableNameMapper.mapName(it) },
38+
schema,
39+
)
40+
41+
override fun transactWriteItems(request: TransactWriteItemsEnhancedRequest): Void? {
42+
return client.transactWriteItems(request)
43+
}
44+
45+
override fun transactWriteItems(requestConsumer: Consumer<TransactWriteItemsEnhancedRequest.Builder>?): Void? {
46+
return client.transactWriteItems(requestConsumer)
47+
}
48+
49+
override fun batchGetItem(request: BatchGetItemEnhancedRequest?): BatchGetResultPageIterable {
50+
return client.batchGetItem(request)
51+
}
52+
53+
override fun batchGetItem(requestConsumer: Consumer<BatchGetItemEnhancedRequest.Builder>?): BatchGetResultPageIterable {
54+
return client.batchGetItem(requestConsumer)
55+
}
56+
57+
override fun batchWriteItem(request: BatchWriteItemEnhancedRequest?): BatchWriteResult {
58+
return client.batchWriteItem(request)
59+
}
60+
61+
override fun batchWriteItem(requestConsumer: Consumer<BatchWriteItemEnhancedRequest.Builder>?): BatchWriteResult {
62+
return client.batchWriteItem(requestConsumer)
63+
}
64+
65+
override fun transactGetItems(request: TransactGetItemsEnhancedRequest?): MutableList<Document> {
66+
return client.transactGetItems(request)
67+
}
68+
69+
override fun transactGetItems(requestConsumer: Consumer<TransactGetItemsEnhancedRequest.Builder>?): MutableList<Document> {
70+
return client.transactGetItems(requestConsumer)
71+
}
72+
73+
override fun transactWriteItemsWithResponse(request: TransactWriteItemsEnhancedRequest?): TransactWriteItemsEnhancedResponse {
74+
return client.transactWriteItemsWithResponse(request)
75+
}
76+
77+
override fun transactWriteItemsWithResponse(requestConsumer: Consumer<TransactWriteItemsEnhancedRequest.Builder>?): TransactWriteItemsEnhancedResponse {
78+
return client.transactWriteItemsWithResponse(requestConsumer)
79+
}
80+
}

‎misk-aws2-dynamodb/src/testFixtures/kotlin/misk/aws2/dynamodb/testing/ExternalTestDynamoDbClientModule.kt‎

Lines changed: 117 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,117 @@
1+
package misk.aws2.dynamodb.testing
2+
3+
import app.cash.tempest2.testing.TestDynamoDbClient
4+
import app.cash.tempest2.testing.TestTable
5+
import app.cash.tempest2.testing.internal.DefaultTestDynamoDbClient
6+
import app.cash.tempest2.testing.internal.createTable
7+
import com.google.common.util.concurrent.AbstractIdleService
8+
import com.google.common.util.concurrent.AbstractService
9+
import com.google.inject.Provides
10+
import jakarta.inject.Inject
11+
import jakarta.inject.Singleton
12+
import misk.ServiceModule
13+
import misk.aws2.dynamodb.DynamoDbService
14+
import misk.aws2.dynamodb.RequiredDynamoDbTable
15+
import misk.aws2.dynamodb.TableNameMapper
16+
import misk.inject.KAbstractModule
17+
import misk.testing.TestFixture
18+
import misk.testing.updateForParallelTests
19+
import software.amazon.awssdk.services.dynamodb.DynamoDbClient
20+
import software.amazon.awssdk.services.dynamodb.model.DeleteTableRequest
21+
import software.amazon.awssdk.services.dynamodb.model.ResourceNotFoundException
22+
import software.amazon.awssdk.services.dynamodb.streams.DynamoDbStreamsClient
23+
24+
/**
25+
* Unlike the InProcessDynamoDbModule and DockerDynamoDbModule classes, this module does not internally start DynamoDB,
26+
* and instead relies on an external DynamoDB server already running on the given port, e.g. using a Gradle task as
27+
* a dependency of the test task, which starts DynamoDB in a Docker container.
28+
*/
29+
class ExternalTestDynamoDbClientModule(
30+
private val port: Int,
31+
originalTables: List<DynamoDbTable>
32+
) : KAbstractModule() {
33+
34+
private val tables = originalTables.map { it.copy(tableName = ParallelTestsTableNameMapper.mapName(it.tableName)) }
35+
36+
constructor(port: Int, vararg tables: DynamoDbTable) : this(port, tables.toList())
37+
38+
override fun configure() {
39+
for (table in tables) {
40+
multibind<DynamoDbTable>().toInstance(table)
41+
}
42+
bind<DynamoDbService>().to<TestDynamoDbService>()
43+
install(ServiceModule<DynamoDbService>().dependsOn<TestDynamoDbFixture>())
44+
install(ServiceModule<TestDynamoDbFixture>())
45+
multibind<TestFixture>().to<TestDynamoDbFixture>()
46+
}
47+
48+
@Provides
49+
@Singleton
50+
fun provideRequiredTables(): List<RequiredDynamoDbTable> = tables.map { RequiredDynamoDbTable(it.tableName) }
51+
52+
@Provides
53+
@Singleton
54+
fun providesTestDynamoDbClient(): TestDynamoDbClient = DefaultTestDynamoDbClient(
55+
tables = tables.map { table ->
56+
TestTable.create(table.tableName, table.tableClass) { table.configureTable(it.toBuilder()).build() }
57+
},
58+
port,
59+
)
60+
61+
@Provides
62+
@Singleton
63+
fun providesAmazonDynamoDB(testDynamoDbClient: TestDynamoDbClient): DynamoDbClient = testDynamoDbClient.dynamoDb
64+
65+
@Provides
66+
@Singleton
67+
fun providesAmazonDynamoDBStreams(testDynamoDbClient: TestDynamoDbClient):
68+
DynamoDbStreamsClient = testDynamoDbClient.dynamoDbStreams
69+
}
70+
71+
@Singleton
72+
private class TestDynamoDbFixture @Inject constructor(
73+
private val client: TestDynamoDbClient,
74+
private val tables: List<DynamoDbTable>
75+
) : AbstractIdleService(), TestFixture {
76+
77+
override fun startUp() {
78+
reset()
79+
}
80+
81+
override fun shutDown() {
82+
}
83+
84+
override fun reset() {
85+
for (tableName in tables.map { it.tableName }) {
86+
try {
87+
client.dynamoDb.deleteTable(DeleteTableRequest.builder().tableName(tableName).build())
88+
} catch (e: ResourceNotFoundException) {
89+
// Ignore if the table doesn't exist
90+
}
91+
}
92+
for (table in tables.map { table ->
93+
TestTable.create(table.tableName, table.tableClass) { table.configureTable(it.toBuilder()).build() }
94+
}) {
95+
client.dynamoDb.createTable(table)
96+
}
97+
}
98+
}
99+
100+
@Singleton
101+
private class TestDynamoDbService @Inject constructor() : AbstractService(), DynamoDbService {
102+
override fun doStart() = notifyStarted()
103+
override fun doStop() = notifyStopped()
104+
}
105+
106+
/**
107+
* A [TableNameMapper] that appends a unique identifier for each test process ID to the table name.
108+
* This is used to ensure that multiple tests can run in parallel without clobbering each other's tables.
109+
*/
110+
object ParallelTestsTableNameMapper : TableNameMapper {
111+
override fun mapName(tableName: String): String {
112+
return tableName.updateForParallelTests { name, index ->
113+
name + "_$index"
114+
}
115+
}
116+
}
117+

0 commit comments

Comments
 (0)