content: package org.aion.db.utils;
  
  import static org.junit.Assert.fail;
  
  import com.spotify.docker.client.DefaultDockerClient;
  import com.spotify.docker.client.DockerClient;
  import com.spotify.docker.client.LogStream;
  import com.spotify.docker.client.messages.ContainerConfig;
  import com.spotify.docker.client.messages.ContainerCreation;
  import com.spotify.docker.client.messages.ExecCreation;
  import com.spotify.docker.client.messages.HostConfig;
  import com.spotify.docker.client.messages.PortBinding;
  import java.util.Arrays;
  import java.util.List;
  import java.util.Map;
  import org.aion.db.impl.DatabaseTestUtils;
  
  /** Helper class for spinning up a MongoDB instance to be used for unit tests. */
  public class MongoTestRunner implements AutoCloseable {
  
      private int port;
      private DockerClient dockerClient;
      private String runningDockerContainerId;
      private static final String MONGO_IMAGE = "library/mongo:3.6.9";
  
      private static class Holder {
          static final MongoTestRunner INSTANCE = new MongoTestRunner();
      }
  
      public static MongoTestRunner inst() {
          return Holder.INSTANCE;
      }
  
      private MongoTestRunner() {
          try {
              // Start by getting a connection to the docker service running on the machine
              DefaultDockerClient.Builder clientBuilder = DefaultDockerClient.fromEnv();
              System.out.println("Connecting to docker daemon at " + clientBuilder.uri().toString());
              dockerClient = clientBuilder.build();
  
              // Pull the docker image, this will be very quick if it already exists on the machine
              dockerClient.pull(
                      MONGO_IMAGE, message -> System.out.println("Docker pull: " + message.status()));
  
              // Bind container port 27017 to an automatically allocated available host port.
              this.port = DatabaseTestUtils.findOpenPort();
              final Map<String, List<PortBinding>> portBindings =
                      Map.of(
                              "27017",
                              Arrays.asList(PortBinding.of("0.0.0.0", Integer.toString(this.port))));
              final HostConfig hostConfig = HostConfig.builder().portBindings(portBindings).build();
  
              // Configure how we want the image to run
              ContainerConfig containerConfig =
                      ContainerConfig.builder()
                              .attachStderr(true)
                              .hostConfig(hostConfig)
                              .exposedPorts("27017")
                              .image(MONGO_IMAGE)
                              .cmd("--replSet", "rs0", "--noauth", "--nojournal", "--quiet")
                              .build();
  
              // Actually start the container
              ContainerCreation creation = dockerClient.createContainer(containerConfig);
              dockerClient.startContainer(creation.id());
              this.runningDockerContainerId = creation.id();
  
              // Next we run a command to initialize the mongo server's replicas set and admin
              // accounts
              String[] initializationCommands = {"mongo", "--eval", "rs.initiate()"};
              tryInitializeDb(initializationCommands, 30, 100);
  
              // Finally, add a shutdown hook to kill the Mongo server when the process dies
              Runtime.getRuntime()
                      .addShutdownHook(
                              new Thread(
                                      () -> {
                                          try {
                                              close();
                                          } catch (Exception e) {
                                              e.printStackTrace();
                                              fail("Failed to close MongoDB connection");
                                          }
                                      }));
  
          } catch (Exception e) {
              e.printStackTrace();
              fail(
                      "Error encountered when initializing mongo docker image. Make sure docker service is running");
          }
      }
  
      /**
       * Helper method to run some initialization command on Mongo with some retry logic if the
       * command fails. Since it's not determinate how long starting the database will take, we need
       * this retry logic.
       *
       * @param initializationCommands The command to actually run
       * @param retriesRemaining How many more times to retry the command if it fails
       * @param pauseTimeMillis How long to pause between retries
       * @throws InterruptedException Thrown when the thread gets interrupted trying to sleep.
       */
      private void tryInitializeDb(
              String[] initializationCommands, int retriesRemaining, long pauseTimeMillis)
              throws InterruptedException {
  
          Exception exception = null;
          String execOutput = "";
          try {
              final ExecCreation execCreation =
                      dockerClient.execCreate(
                              this.runningDockerContainerId,
                              initializationCommands,
                              DockerClient.ExecCreateParam.attachStdout(),
                              DockerClient.ExecCreateParam.attachStderr(),
                              DockerClient.ExecCreateParam.detach(false));
              final LogStream output = dockerClient.execStart(execCreation.id());
              execOutput = output.readFully();
          } catch (Exception e) {
              exception = e;
          }
  
          // We can't get the exit code, but look for an expected message in the output to determine
          // success
          if (exception != null
                  || !execOutput.contains("Using a default configuration for the set")) {
              // This is the case that the command didn't work
              if (retriesRemaining == 0) {
                  // We're out of retries, we should fail
                  if (exception != null) {
                      exception.printStackTrace();
                  }
  
                  fail(
                          "Failed to initialize MongoDB, no retries remaining. Output was: "
                                  + execOutput);
              } else {
                  Thread.sleep(pauseTimeMillis);
                  tryInitializeDb(initializationCommands, retriesRemaining - 1, pauseTimeMillis);
              }
          }
      }
  
      /**
       * Returns the connection string to be used to connect to the started Mongo instance
       *
       * @return The connection string.
       */
      public String getConnectionString() {
          return String.format("mongodb://localhost:%d", this.port);
      }
  
      @Override
      public void close() throws Exception {
          if (this.dockerClient != null && this.runningDockerContainerId != null) {
              System.out.println("Killing mongo docker container");
              this.dockerClient.killContainer(this.runningDockerContainerId);
              this.dockerClient.removeContainer(this.runningDockerContainerId);
              this.dockerClient.close();
              this.dockerClient = null;
              this.runningDockerContainerId = null;
          }
      }
  }
  