created db provisioner for programmatic provisioning

Author: jruaux

src/main/java/com/redislabs/testcontainers/RedisClusterContainer.java

@@ -1,10 +1,10 @@
 package com.redislabs.testcontainers;
 
-import org.testcontainers.containers.FixedHostPortGenericContainer;
+import org.testcontainers.containers.GenericContainer;
 import org.testcontainers.containers.wait.strategy.Wait;
 import org.testcontainers.shaded.org.apache.commons.lang.ClassUtils;
 
-public class RedisClusterContainer extends FixedHostPortGenericContainer<RedisClusterContainer> implements RedisServer {
+public class RedisClusterContainer extends GenericContainer<RedisClusterContainer> implements RedisServer {
 
     private static final String ENV_MASTERS = "MASTERS";
     private static final String ENV_SLAVES_PER_MASTER = "SLAVES_PER_MASTER";
@@ -54,7 +54,7 @@ private RedisClusterContainer update() {
         withEnv(ENV_MASTERS, String.valueOf(masters));
         withEnv(ENV_SLAVES_PER_MASTER, String.valueOf(slavesPerMaster));
         for (int port : ports()) {
-            withFixedExposedPort(port, port);
+            addFixedExposedPort(port, port);
         }
         return this;
     }

src/main/java/com/redislabs/testcontainers/RedisEnterpriseContainer.java

@@ -5,98 +5,69 @@
 import com.github.dockerjava.api.command.InspectContainerResponse;
 import com.github.dockerjava.api.exception.DockerException;
 import com.github.dockerjava.zerodep.shaded.org.apache.hc.client5.http.auth.UsernamePasswordCredentials;
-import com.redislabs.testcontainers.support.enterprise.ActionStatus;
-import com.redislabs.testcontainers.support.enterprise.Command;
-import com.redislabs.testcontainers.support.enterprise.CommandResponse;
-import com.redislabs.testcontainers.support.enterprise.Database;
-import com.redislabs.testcontainers.support.enterprise.DatabaseCreateResponse;
-import com.redislabs.testcontainers.support.enterprise.Module;
-import com.redislabs.testcontainers.support.enterprise.ModuleInstallResponse;
+import com.redislabs.testcontainers.support.enterprise.DatabaseProvisioner;
 import com.redislabs.testcontainers.support.enterprise.RestAPI;
+import com.redislabs.testcontainers.support.enterprise.rest.Database;
+import com.redislabs.testcontainers.support.enterprise.rest.DatabaseCreateResponse;
 import lombok.AccessLevel;
 import lombok.AllArgsConstructor;
-import lombok.Getter;
 import lombok.SneakyThrows;
 import lombok.Value;
 import lombok.extern.slf4j.Slf4j;
 import org.testcontainers.DockerClientFactory;
-import org.testcontainers.containers.ContainerLaunchException;
 import org.testcontainers.containers.GenericContainer;
 import org.testcontainers.containers.output.FrameConsumerResultCallback;
 import org.testcontainers.containers.output.OutputFrame;
 import org.testcontainers.containers.output.ToStringConsumer;
 import org.testcontainers.containers.wait.strategy.Wait;
-import org.testcontainers.shaded.org.apache.commons.io.IOUtils;
+import org.testcontainers.shaded.org.apache.commons.lang.ClassUtils;
 import org.testcontainers.utility.DockerImageName;
 import org.testcontainers.utility.TestEnvironment;
 
-import java.io.ByteArrayOutputStream;
 import java.io.IOException;
-import java.io.InputStream;
-import java.net.URL;
 import java.nio.charset.StandardCharsets;
 import java.time.Duration;
-import java.util.Arrays;
-import java.util.Collections;
-import java.util.HashSet;
-import java.util.List;
-import java.util.Map;
-import java.util.Set;
-import java.util.function.Function;
-import java.util.stream.Collectors;
-import java.util.stream.Stream;
-
-@SuppressWarnings({"unused", "UnusedReturnValue", "BusyWait"})
+
+@SuppressWarnings({"unused", "UnusedReturnValue"})
 @Slf4j
 public class RedisEnterpriseContainer extends GenericContainer<RedisEnterpriseContainer> implements RedisServer {
 
+    public static final String MODULE_BLOOM = "bf";
+    public static final String MODULE_GEARS = "rg";
+    public static final String MODULE_GRAPH = "graph";
+    public static final String MODULE_JSON = "ReJSON";
+    public static final String MODULE_SEARCH = "search";
+    public static final String MODULE_TIMESERIES = "timeseries";
+
     public static final String ADMIN_USERNAME = "testcontainers@redislabs.com";
     public static final String ADMIN_PASSWORD = "redislabs123";
-    public static final String GEARS_MODULE_FILE = "redisgears.linux-bionic-x64.1.0.6.zip";
-    public static final String URL_GEARS_MODULE = "https://redismodules.s3.amazonaws.com/redisgears/" + GEARS_MODULE_FILE;
     public static final DockerImageName DEFAULT_IMAGE_NAME = DockerImageName.parse("redislabs/redis");
-    public static final List<Database.Regex> DEFAULT_SHARD_KEY_REGEX = Arrays.asList(Database.Regex.of(".*\\{(?<tag>.*)\\}.*"), Database.Regex.of("(?<tag>.*)"));
-    private static final long DEFAULT_MODULE_INSTALLATION_TIMEOUT = 60000;
-    private static final long DEFAULT_MODULE_INSTALLATION_CHECK_INTERVAL = 5000;
-    private static final long DEFAULT_DATABASE_PING_TIMEOUT = 10000;
-    private static final long DEFAULT_DATABASE_PING_RETRY_INTERVAL = 1000;
     private static final long DEFAULT_RLADMIN_WAIT_INTERVAL = 500;
     private static final String NODE_EXTERNAL_ADDR = "0.0.0.0";
-    private static final int DEFAULT_CLUSTER_SHARD_COUNT = 3;
+    private static final int DEFAULT_SHARD_COUNT = 2;
+    private static final String DEFAULT_DATABASE_NAME = "testcontainers";
     public static int ADMIN_PORT = 8443;
     public static int ENDPOINT_PORT = 12000;
-    private static final Map<String, RedisModule> MODULES = Stream.of(RedisModule.values()).collect(Collectors.toMap(RedisModule::getModuleName, Function.identity()));
-
-    private final Set<RedisModule> modules = new HashSet<>();
-    private RestAPI restAPI = RestAPI.credentials(new UsernamePasswordCredentials(ADMIN_USERNAME, ADMIN_PASSWORD.toCharArray())).build();
-    private Duration moduleInstallationTimeout = Duration.ofMillis(DEFAULT_MODULE_INSTALLATION_TIMEOUT);
-    private Duration moduleInstallationCheckInterval = Duration.ofMillis(DEFAULT_MODULE_INSTALLATION_CHECK_INTERVAL);
-    private Duration databasePingTimeout = Duration.ofMillis(DEFAULT_DATABASE_PING_TIMEOUT);
-    private Database database = Database.builder().name("testcontainers").build();
-    private Duration databasePingRetryInterval = Duration.ofMillis(DEFAULT_DATABASE_PING_RETRY_INTERVAL);
+
+    private DatabaseProvisioner.Options provisionerOptions = DatabaseProvisioner.Options.builder().build();
     private Duration rladminWaitDuration = Duration.ofMillis(DEFAULT_RLADMIN_WAIT_INTERVAL);
+    private String databaseName = DEFAULT_DATABASE_NAME;
+    private int shardCount = DEFAULT_SHARD_COUNT;
+    private boolean ossCluster;
+    private Database.Module[] modules = new Database.Module[0];
 
     public RedisEnterpriseContainer() {
         super(DEFAULT_IMAGE_NAME);
+        addFixedExposedPort(RestAPI.DEFAULT_PORT, RestAPI.DEFAULT_PORT);
         addFixedExposedPort(ADMIN_PORT, ADMIN_PORT);
         addFixedExposedPort(ENDPOINT_PORT, ENDPOINT_PORT);
         withPrivilegedMode(true);
         withPublishAllPorts(false);
         waitingFor(Wait.forLogMessage(".*success: job_scheduler entered RUNNING state, process has stayed up for.*\\n", 1));
     }
 
-    public RedisEnterpriseContainer withRestAPI(RestAPI restAPI) {
-        this.restAPI = restAPI;
-        return this;
-    }
-
-    public RedisEnterpriseContainer withDatabase(Database database) {
-        this.database = database;
-        return this;
-    }
-
-    public RedisEnterpriseContainer withModuleInstallationCheckInterval(Duration moduleInstallationCheckInterval) {
-        this.moduleInstallationCheckInterval = moduleInstallationCheckInterval;
+    public RedisEnterpriseContainer withProvisionerOptions(DatabaseProvisioner.Options options) {
+        this.provisionerOptions = options;
         return this;
     }
 
@@ -106,109 +77,40 @@ public RedisEnterpriseContainer withRladminWaitDuration(Duration rladminWaitDura
     }
 
     public RedisEnterpriseContainer withShardCount(int shardCount) {
-        this.database.setShardCount(shardCount);
-        if (shardCount > 1) {
-            this.database.setSharding(true);
-            this.database.setShardKeyRegex(DEFAULT_SHARD_KEY_REGEX);
-        }
+        this.shardCount = shardCount;
         return this;
     }
 
-    public RedisEnterpriseContainer withModuleInstallationTimeout(Duration moduleInstallationTimeout) {
-        this.moduleInstallationTimeout = moduleInstallationTimeout;
-        return this;
-    }
-
-    public RedisEnterpriseContainer withDatabasePingTimeout(Duration databasePingTimeout) {
-        this.databasePingTimeout = databasePingTimeout;
-        return this;
-    }
-
-    public RedisEnterpriseContainer databasePingRetryInterval(Duration databasePingRetryInterval) {
-        this.databasePingRetryInterval = databasePingRetryInterval;
-        return this;
-    }
-
-    public RedisEnterpriseContainer withOSSCluster() {
-        this.database.setOssCluster(true);
-        this.database.setProxyPolicy(Database.ProxyPolicy.ALL_MASTER_SHARDS);
-        this.database.setOssClusterAPIPreferredIPType(Database.IPType.EXTERNAL);
-        if (database.getShardCount() == null || database.getShardCount() < 2) {
-            log.info("Setting shard count to {}", DEFAULT_CLUSTER_SHARD_COUNT);
-            withShardCount(DEFAULT_CLUSTER_SHARD_COUNT);
-        }
+    public RedisEnterpriseContainer withModules(Database.Module... modules) {
+        this.modules = modules;
         return this;
     }
 
-    public RedisEnterpriseContainer withModules(RedisModule... modules) {
-        Collections.addAll(this.modules, modules);
+    public RedisEnterpriseContainer withDatabaseName(String name) {
+        this.databaseName = name;
         return this;
     }
 
-    @Override
-    public void start() {
-        addFixedExposedPort(restAPI.getPort(), restAPI.getPort());
-        super.start();
-    }
-
     @SneakyThrows
     @Override
     protected void containerIsStarted(InspectContainerResponse containerInfo) {
         super.containerIsStarted(containerInfo);
+        log.info("Creating Redis Enterprise cluster");
         execute("/opt/redislabs/bin/rladmin", "cluster", "create", "name", "cluster.local", "username", ADMIN_USERNAME, "password", ADMIN_PASSWORD);
         Thread.sleep(rladminWaitDuration.toMillis());
+        log.info("Setting Redis Enterprise node external IP");
         execute("/opt/redislabs/bin/rladmin", "node", "1", "external_addr", "set", NODE_EXTERNAL_ADDR);
         Thread.sleep(rladminWaitDuration.toMillis());
-        if (!modules.isEmpty()) {
-            if (modules.contains(RedisModule.GEARS)) {
-                installGears();
-            }
-            List<Module> availableModules = restAPI.modules();
-            Map<RedisModule, String> moduleIds = availableModules.stream().collect(Collectors.toMap(m -> MODULES.get(m.getName()), Module::getId));
-            database.setModuleList(modules.stream().map(m -> Database.Module.builder().id(moduleIds.get(m)).name(m.getModuleName()).build()).collect(Collectors.toList()));
-        }
-        DatabaseCreateResponse databaseCreateResponse = restAPI.create(database);
-        if (canPingDatabase(databaseCreateResponse.getUid())) {
-            log.info("Created database {} with UID {}", databaseCreateResponse.getName(), databaseCreateResponse.getUid());
-        } else {
-            throw new ContainerLaunchException("Could not ping database at " + getRedisURI());
-        }
+        RestAPI restAPI = RestAPI.credentials(new UsernamePasswordCredentials(ADMIN_USERNAME, ADMIN_PASSWORD.toCharArray())).build();
+        DatabaseProvisioner provisioner = DatabaseProvisioner.restAPI(restAPI).options(provisionerOptions).build();
+        Database database = Database.name(databaseName).port(ENDPOINT_PORT).shardCount(shardCount).ossCluster(ossCluster).modules(modules).build();
+        DatabaseCreateResponse response = provisioner.create(database);
+        log.info("Created database {} with UID {}", database.getName(), response.getUid());
     }
 
-
-    private boolean canPingDatabase(long uid) throws Exception {
-        long start = System.currentTimeMillis();
-        do {
-            try {
-                CommandResponse response = restAPI.command(uid, Command.command("PING").build());
-                if (response.getResponse().asBoolean()) {
-                    return true;
-                }
-            } catch (Exception e) {
-                // ignore
-            }
-            Thread.sleep(databasePingRetryInterval.toMillis());
-        } while (System.currentTimeMillis() - start < databasePingTimeout.toMillis());
-        return false;
-    }
-
-    private void installGears() throws Exception {
-        try (InputStream zipInputStream = new URL(URL_GEARS_MODULE).openStream()) {
-            ByteArrayOutputStream baos = new ByteArrayOutputStream();
-            IOUtils.copy(zipInputStream, baos);
-            ModuleInstallResponse installResponse = restAPI.module(GEARS_MODULE_FILE, baos.toByteArray());
-            long start = System.currentTimeMillis();
-            ActionStatus status;
-            do {
-                status = restAPI.actionStatus(installResponse.getActionUID());
-                if ("completed".equals(status.getStatus())) {
-                    return;
-                }
-                // Wait before checking again
-                Thread.sleep(moduleInstallationCheckInterval.toMillis());
-            } while (System.currentTimeMillis() - start < moduleInstallationTimeout.toMillis());
-            throw new ContainerLaunchException("Timed out waiting for module installation to complete. Action UID: " + installResponse.getActionUID());
-        }
+    public RedisEnterpriseContainer withOSSCluster() {
+        this.ossCluster = true;
+        return this;
     }
 
     @Override
@@ -218,7 +120,11 @@ public String getRedisURI() {
 
     @Override
     public boolean isCluster() {
-        return Boolean.TRUE.equals(database.getOssCluster());
+        return ossCluster;
+    }
+
+    public String toString() {
+        return ClassUtils.getShortClassName(this.getClass());
     }
 
     @SuppressWarnings("deprecation")
@@ -278,18 +184,6 @@ private boolean isRunning(InspectContainerResponse containerInfo) {
         }
     }
 
-    public enum RedisModule {
-
-        JSON("ReJSON"), SEARCH("search"), TIMESERIES("timeseries"), BLOOM("bf"), GRAPH("graph"), GEARS("rg");
-
-        @Getter
-        private final String moduleName;
-
-        RedisModule(String name) {
-            this.moduleName = name;
-        }
-
-    }
 
 }