Migrate from embedded DynamoDBLocal to Testcontainers

service/src/test/java/org/whispersystems/textsecuregcm/WhisperServerServiceTest.java

@@ -9,6 +9,7 @@ import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.junit.jupiter.api.Assertions.assertFalse;
 import static org.junit.jupiter.api.Assertions.assertTrue;
 
+import io.dropwizard.testing.ConfigOverride;
 import io.dropwizard.testing.junit5.DropwizardAppExtension;
 import io.dropwizard.testing.junit5.DropwizardExtensionsSupport;
 import io.dropwizard.util.Resources;
@@ -25,6 +26,7 @@ import org.junit.jupiter.api.AfterAll;
 import org.junit.jupiter.api.BeforeAll;
 import org.junit.jupiter.api.Test;
 import org.junit.jupiter.api.extension.ExtendWith;
+import org.junit.jupiter.api.extension.RegisterExtension;
 import org.whispersystems.textsecuregcm.metrics.NoopAwsSdkMetricPublisher;
 import org.whispersystems.textsecuregcm.storage.DynamoDbExtension;
 import org.whispersystems.textsecuregcm.storage.DynamoDbExtensionSchema;
@@ -51,8 +53,12 @@ class WhisperServerServiceTest {
   private static final WebSocketClient webSocketClient = new WebSocketClient();
 
   private static final DropwizardAppExtension<WhisperServerConfiguration> EXTENSION = new DropwizardAppExtension<>(
-      WhisperServerService.class, Resources.getResource("config/test.yml").getPath());
+      WhisperServerService.class, Resources.getResource("config/test.yml").getPath(),
+      // Tables will be created by the local DynamoDbExtension
+      ConfigOverride.config("dynamoDbClient.initTables", "false"));
 
+  @RegisterExtension
+  public static final DynamoDbExtension DYNAMO_DB_EXTENSION = new DynamoDbExtension(DynamoDbExtensionSchema.Tables.values());
 
   @AfterAll
   static void teardown() {
@@ -118,8 +124,6 @@ class WhisperServerServiceTest {
     assertEquals(401, whoami.getStatus());
     final long whoamiTimestamp = Long.parseLong(whoami.getHeaders().get(HeaderUtils.TIMESTAMP_HEADER.toLowerCase()));
     assertTrue(whoamiTimestamp >= start);
-
-
   }
 
   @Test
@@ -142,11 +146,7 @@ class WhisperServerServiceTest {
   void dynamoDb() {
     // confirm that local dynamodb nominally works
 
-    final AwsCredentialsProvider awsCredentialsProvider = EXTENSION.getConfiguration().getAwsCredentialsConfiguration()
-        .build();
-
-    final DynamoDbClient dynamoDbClient = EXTENSION.getConfiguration().getDynamoDbClientConfiguration()
-        .buildSyncClient(awsCredentialsProvider, new NoopAwsSdkMetricPublisher());
+    final DynamoDbClient dynamoDbClient = getDynamoDbClient();
 
     final DynamoDbExtension.TableSchema numbers = DynamoDbExtensionSchema.Tables.NUMBERS;
     final AttributeValue numberAV = AttributeValues.s("+12125550001");
@@ -176,4 +176,12 @@ class WhisperServerServiceTest {
         .build());
   }
 
+  private static DynamoDbClient getDynamoDbClient() {
+    final AwsCredentialsProvider awsCredentialsProvider = EXTENSION.getConfiguration().getAwsCredentialsConfiguration()
+        .build();
+
+    return EXTENSION.getConfiguration().getDynamoDbClientConfiguration()
+        .buildSyncClient(awsCredentialsProvider, new NoopAwsSdkMetricPublisher());
+  }
+
 }

service/src/test/java/org/whispersystems/textsecuregcm/configuration/LocalDynamoDbFactory.java

@@ -5,6 +5,7 @@
 
 package org.whispersystems.textsecuregcm.configuration;
 
+import com.fasterxml.jackson.annotation.JsonProperty;
 import com.fasterxml.jackson.annotation.JsonTypeName;
 import org.whispersystems.textsecuregcm.storage.DynamoDbExtension;
 import org.whispersystems.textsecuregcm.storage.DynamoDbExtensionSchema;
@@ -18,23 +19,50 @@ public class LocalDynamoDbFactory implements DynamoDbClientFactory {
 
   private static final DynamoDbExtension EXTENSION = new DynamoDbExtension(DynamoDbExtensionSchema.Tables.values());
 
-  static {
+  /**
+   * If true, tables will be created the first time a DynamoDB client is built.
+   * <p>
+   * Defaults to {@code true}.
+   */
+  @JsonProperty
+  boolean initTables = true;
+
+  public LocalDynamoDbFactory() {
     try {
-      EXTENSION.beforeEach(null);
+      EXTENSION.beforeAll(null);
     } catch (Exception e) {
       throw new RuntimeException(e);
     }
 
-    Runtime.getRuntime().addShutdownHook(new Thread(() -> EXTENSION.afterEach(null)));
+    Runtime.getRuntime().addShutdownHook(new Thread(() -> {
+      try {
+        EXTENSION.close();
+      } catch (Throwable e) {
+        throw new RuntimeException(e);
+      }
+    }));
   }
 
   @Override
   public DynamoDbClient buildSyncClient(final AwsCredentialsProvider awsCredentialsProvider, final MetricPublisher metricPublisher) {
+    initTablesIfNecessary();
     return EXTENSION.getDynamoDbClient();
   }
 
   @Override
   public DynamoDbAsyncClient buildAsyncClient(final AwsCredentialsProvider awsCredentialsProvider, final MetricPublisher metricPublisher) {
+    initTablesIfNecessary();
     return EXTENSION.getDynamoDbAsyncClient();
   }
+
+  private void initTablesIfNecessary() {
+    try {
+      if (initTables) {
+        EXTENSION.beforeEach(null);
+        initTables = false;
+      }
+    } catch (Exception e) {
+      throw new RuntimeException(e);
+    }
+  }
 }

service/src/test/java/org/whispersystems/textsecuregcm/redis/RedisClusterExtension.java

@@ -31,6 +31,7 @@ import org.testcontainers.containers.wait.strategy.Wait;
 import org.testcontainers.containers.wait.strategy.WaitStrategy;
 import org.whispersystems.textsecuregcm.configuration.CircuitBreakerConfiguration;
 import org.whispersystems.textsecuregcm.configuration.RetryConfiguration;
+import org.whispersystems.textsecuregcm.util.TestcontainersImages;
 
 public class RedisClusterExtension implements BeforeAllCallback, BeforeEachCallback, AfterEachCallback, ExtensionContext.Store.CloseableResource {
 
@@ -52,24 +53,22 @@ public class RedisClusterExtension implements BeforeAllCallback, BeforeEachCallb
 
   private static final String[] REDIS_SERVICE_NAMES = new String[] { "redis-0-1", "redis-1-1", "redis-2-1" };
 
-  // The image we're using is bitnami/redis-cluster:7.4; please see
-  // https://hub.docker.com/layers/bitnami/redis-cluster/7.4/images/sha256-c11efe6a53692829b6e031ea8b5b4caa380df3c84ad4242549851d345592708d
-  private static final String CLUSTER_COMPOSE_FILE_CONTENTS = """
+  private static final String CLUSTER_COMPOSE_FILE_CONTENTS = String.format("""
       services:
         redis-0:
-          image: docker.io/bitnami/redis-cluster@sha256:a53d023fdfaf8a8d7ddc58da040d3494e4cb45772644618ffa44c42dcd32b9af
+          image: %1$s
           environment:
             - 'ALLOW_EMPTY_PASSWORD=yes'
             - 'REDIS_NODES=redis-0 redis-1 redis-2'
 
         redis-1:
-          image: docker.io/bitnami/redis-cluster@sha256:a53d023fdfaf8a8d7ddc58da040d3494e4cb45772644618ffa44c42dcd32b9af
+          image: %1$s
           environment:
             - 'ALLOW_EMPTY_PASSWORD=yes'
             - 'REDIS_NODES=redis-0 redis-1 redis-2'
 
         redis-2:
-          image: docker.io/bitnami/redis-cluster@sha256:a53d023fdfaf8a8d7ddc58da040d3494e4cb45772644618ffa44c42dcd32b9af
+          image: %1$s
           depends_on:
             - redis-0
             - redis-1
@@ -78,7 +77,7 @@ public class RedisClusterExtension implements BeforeAllCallback, BeforeEachCallb
             - 'REDIS_CLUSTER_REPLICAS=0'
             - 'REDIS_NODES=redis-0 redis-1 redis-2'
             - 'REDIS_CLUSTER_CREATOR=yes'
-      """;
+      """, TestcontainersImages.getRedisCluster());
 
   public RedisClusterExtension(final Duration timeout, final RetryConfiguration retryConfiguration) {
     this.timeout = timeout;

service/src/test/java/org/whispersystems/textsecuregcm/redis/RedisServerExtension.java

@@ -17,6 +17,7 @@ import org.junit.jupiter.api.extension.ExtensionContext;
 import org.testcontainers.utility.DockerImageName;
 import org.whispersystems.textsecuregcm.configuration.CircuitBreakerConfiguration;
 import org.whispersystems.textsecuregcm.configuration.RetryConfiguration;
+import org.whispersystems.textsecuregcm.util.TestcontainersImages;
 
 public class RedisServerExtension implements BeforeAllCallback, BeforeEachCallback, AfterEachCallback, ExtensionContext.Store.CloseableResource {
 
@@ -25,8 +26,7 @@ public class RedisServerExtension implements BeforeAllCallback, BeforeEachCallba
   private ClientResources redisClientResources;
   private FaultTolerantRedisClient faultTolerantRedisClient;
 
-  // redis:7.4-apline; see https://hub.docker.com/layers/library/redis/7.4-alpine/images/sha256-e1b05db81cda983ede3bbb3e834e7ebec8faafa275f55f7f91f3ee84114f98a7
-  private static final DockerImageName REDIS_IMAGE = DockerImageName.parse("redis@sha256:af1d0fc3f63b02b13ff7906c9baf7c5b390b8881ca08119cd570677fe2f60b55");
+  private static final DockerImageName REDIS_IMAGE = DockerImageName.parse(TestcontainersImages.getRedis());
 
   public static class RedisServerExtensionBuilder {
 

service/src/test/java/org/whispersystems/textsecuregcm/storage/DynamoDbExtension.java

@@ -5,12 +5,21 @@
 
 package org.whispersystems.textsecuregcm.storage;
 
-import com.amazonaws.services.dynamodbv2.local.embedded.DynamoDBEmbedded;
-import com.amazonaws.services.dynamodbv2.local.shared.access.AmazonDynamoDBLocal;
+import java.net.URI;
+import java.time.Duration;
+import java.time.Instant;
 import java.util.List;
+import org.junit.jupiter.api.extension.AfterAllCallback;
 import org.junit.jupiter.api.extension.AfterEachCallback;
+import org.junit.jupiter.api.extension.BeforeAllCallback;
 import org.junit.jupiter.api.extension.BeforeEachCallback;
 import org.junit.jupiter.api.extension.ExtensionContext;
+import org.testcontainers.containers.GenericContainer;
+import org.testcontainers.utility.DockerImageName;
+import org.whispersystems.textsecuregcm.util.TestcontainersImages;
+import software.amazon.awssdk.auth.credentials.AwsBasicCredentials;
+import software.amazon.awssdk.auth.credentials.StaticCredentialsProvider;
+import software.amazon.awssdk.regions.Region;
 import software.amazon.awssdk.services.dynamodb.DynamoDbAsyncClient;
 import software.amazon.awssdk.services.dynamodb.DynamoDbClient;
 import software.amazon.awssdk.services.dynamodb.model.AttributeDefinition;
@@ -20,8 +29,9 @@ import software.amazon.awssdk.services.dynamodb.model.KeySchemaElement;
 import software.amazon.awssdk.services.dynamodb.model.KeyType;
 import software.amazon.awssdk.services.dynamodb.model.LocalSecondaryIndex;
 import software.amazon.awssdk.services.dynamodb.model.ProvisionedThroughput;
+import software.amazon.awssdk.services.dynamodb.model.ResourceNotFoundException;
 
-public class DynamoDbExtension implements BeforeEachCallback, AfterEachCallback {
+public class DynamoDbExtension implements BeforeAllCallback, BeforeEachCallback, AfterEachCallback, AfterAllCallback, ExtensionContext.Store.CloseableResource {
 
   public interface TableSchema {
     String tableName();
@@ -46,40 +56,104 @@ public class DynamoDbExtension implements BeforeEachCallback, AfterEachCallback
       .writeCapacityUnits(20L)
       .build();
 
-  private AmazonDynamoDBLocal embedded;
+  private static final DockerImageName DYNAMO_DB_IMAGE = DockerImageName.parse(TestcontainersImages.getDynamoDb());
+  private static final int CONTAINER_PORT = 8000;
+  private static final GenericContainer<?> dynamoDbContainer =  new GenericContainer<>(DYNAMO_DB_IMAGE)
+      .withExposedPorts(CONTAINER_PORT)
+      .withCommand("-jar DynamoDBLocal.jar -inMemory -sharedDb -disableTelemetry");
+
   private final List<TableSchema> schemas;
-  private DynamoDbClient dynamoDB2;
-  private DynamoDbAsyncClient dynamoAsyncDB2;
+  private DynamoDbClient dynamoDb;
+  private DynamoDbAsyncClient dynamoDbAsync;
 
   public DynamoDbExtension(TableSchema... schemas) {
     this.schemas = List.of(schemas);
   }
 
+  /**
+   * Starts the DynamoDB server
+   */
   @Override
-  public void afterEach(ExtensionContext context) {
-    stopServer();
+  public void beforeAll(ExtensionContext context) throws Exception {
+    startServer();
   }
 
   /**
-   * For use in integration tests that want to test resiliency/error handling
+   * Creates the tables from {@link #schemas}
    */
-  public void stopServer() {
+  @Override
+  public void beforeEach(final ExtensionContext context) throws Exception {
+    createTables();
+  }
+
+  /**
+   * Deletes the tables from {@link #schemas}
+   */
+  @Override
+  public void afterEach(ExtensionContext context) {
+    final Instant timeout = Instant.now().plus(Duration.ofSeconds(1));
+
+    schemas.stream().map(tableSchema -> dynamoDb.deleteTable(builder -> builder.tableName(tableSchema.tableName())))
+        .forEach(deleteTableResponse -> {
+          while (Instant.now().isBefore(timeout)) {
+            try {
+              // `deleteTable` is technically asynchronous, although it seems to be uncommon with DynamoDB Local,
+              // so this will usually throw and very rarely sleep().
+              dynamoDb.describeTable(builder -> builder.tableName(deleteTableResponse.tableDescription().tableName()));
+              Thread.sleep(50);
+            } catch (ResourceNotFoundException ignored) {
+              // success
+              break;
+            } catch (InterruptedException e) {
+              throw new RuntimeException(e);
+            }
+          }
+        });
+  }
+
+  @Override
+  public void afterAll(ExtensionContext context) throws Exception {
+    dynamoDb.close();
+    dynamoDbAsync.close();
+  }
+
+  @Override
+  public void close() throws Throwable {
+    stopServer();
+  }
+
+  private void startServer() {
+    dynamoDbContainer.start();
+    initializeClient();
+  }
+
+  private void stopServer() {
     try {
-      embedded.shutdown();
+      if (dynamoDbContainer != null) {
+        dynamoDb.close();
+        dynamoDb = null;
+
+        dynamoDbAsync.close();
+        dynamoDbAsync = null;
+
+        dynamoDbContainer.stop();
+      }
     } catch (Exception e) {
       throw new RuntimeException(e);
     }
   }
 
-  @Override
-  public void beforeEach(ExtensionContext context) throws Exception {
-    initializeClient();
-
+  /**
+   * For use in integration tests that want to test resiliency/error handling
+   */
+  public void resetServer() {
+    stopServer();
+    startServer();
     createTables();
   }
 
   private void createTables() {
-    schemas.stream().forEach(this::createTable);
+    schemas.forEach(this::createTable);
   }
 
   private void createTable(TableSchema schema) {
@@ -108,17 +182,27 @@ public class DynamoDbExtension implements BeforeEachCallback, AfterEachCallback
   }
 
   private void initializeClient() {
-    embedded = DynamoDBEmbedded.create();
-    dynamoDB2 = embedded.dynamoDbClient();
-    dynamoAsyncDB2 = embedded.dynamoDbAsyncClient();
+    final URI endpoint = URI.create(
+        String.format("http://%s:%d", dynamoDbContainer.getHost(), dynamoDbContainer.getMappedPort(CONTAINER_PORT)));
+
+    dynamoDb = DynamoDbClient.builder()
+        .region(Region.of("local"))
+        .credentialsProvider(StaticCredentialsProvider.create(AwsBasicCredentials.create("test", "test")))
+        .endpointOverride(endpoint)
+        .build();
+    dynamoDbAsync = DynamoDbAsyncClient.builder()
+        .region(Region.of("local"))
+        .credentialsProvider(StaticCredentialsProvider.create(AwsBasicCredentials.create("test", "test")))
+        .endpointOverride(endpoint)
+        .build();
   }
 
   public DynamoDbClient getDynamoDbClient() {
-    return dynamoDB2;
+    return dynamoDb;
   }
 
   public DynamoDbAsyncClient getDynamoDbAsyncClient() {
-    return dynamoAsyncDB2;
+    return dynamoDbAsync;
   }
 
 }

service/src/test/java/org/whispersystems/textsecuregcm/storage/MessagesDynamoDbTest.java

@@ -177,7 +177,7 @@ class MessagesDynamoDbTest {
         .thenRequest(halfOfMessageLoadLimit)
         .expectNextCount(halfOfMessageLoadLimit)
         // the first 100 should be fetched and buffered, but further requests should fail
-        .then(DYNAMO_DB_EXTENSION::stopServer)
+        .then(DYNAMO_DB_EXTENSION::resetServer)
         .thenRequest(halfOfMessageLoadLimit)
         .expectNextCount(halfOfMessageLoadLimit)
         // we’ve consumed all the buffered messages, so a single request will fail

service/src/test/java/org/whispersystems/textsecuregcm/storage/S3LocalStackExtension.java

@@ -7,7 +7,6 @@ package org.whispersystems.textsecuregcm.storage;
 
 import static org.testcontainers.containers.localstack.LocalStackContainer.Service.S3;
 
-import java.util.Objects;
 import org.junit.jupiter.api.extension.AfterAllCallback;
 import org.junit.jupiter.api.extension.AfterEachCallback;
 import org.junit.jupiter.api.extension.BeforeAllCallback;
@@ -16,6 +15,7 @@ import org.junit.jupiter.api.extension.ExtensionContext;
 import org.testcontainers.containers.localstack.LocalStackContainer;
 import org.testcontainers.junit.jupiter.Testcontainers;
 import org.testcontainers.utility.DockerImageName;
+import org.whispersystems.textsecuregcm.util.TestcontainersImages;
 import reactor.core.publisher.Flux;
 import reactor.core.publisher.Mono;
 import software.amazon.awssdk.auth.credentials.AwsBasicCredentials;
@@ -31,10 +31,7 @@ import software.amazon.awssdk.services.s3.model.ListObjectsV2Request;
 public class S3LocalStackExtension implements BeforeEachCallback, AfterEachCallback, BeforeAllCallback,
     AfterAllCallback {
 
-  private final static DockerImageName LOCAL_STACK_IMAGE =
-      DockerImageName.parse(Objects.requireNonNull(
-          System.getProperty("localstackImage"),
-          "Local stack image not found; must provide localstackImage system property"));
+  private final static DockerImageName LOCAL_STACK_IMAGE = DockerImageName.parse(TestcontainersImages.getLocalStack());
 
   private static LocalStackContainer LOCAL_STACK = new LocalStackContainer(LOCAL_STACK_IMAGE).withServices(S3);