FaunaDB container testing trait for Scala

FaunaDBTesting.scala

import java.net.{ InetSocketAddress, Socket }
import java.time.Duration
import java.time.temporal.ChronoUnit.SECONDS
import java.util.concurrent.{ TimeUnit, TimeoutException }
import java.util.function.{ Consumer, Predicate }

import com.dimafeng.testcontainers.{ Container, GenericContainer }
import com.faunadb.client.FaunaClient
import org.junit.runner.Description
import org.slf4j.LoggerFactory
import org.testcontainers.DockerClientFactory
import org.testcontainers.containers.ContainerLaunchException
import org.testcontainers.containers.output.{ OutputFrame, Slf4jLogConsumer, WaitingConsumer }
import org.testcontainers.containers.wait.strategy.AbstractWaitStrategy
import org.testcontainers.utility.LogUtils

import scala.util.{ Failure, Success, Try }

/**
  * A mix-in trait, which can be used to spool up FaunaDB container in tests
  */
trait FaunaDBTesting {

  /**
    * Attempt to connect to 127.0.0.1:8443 before trying to spool up a container,
    * so one could speed up local testing if Fauna DB is already running
    */
  val faunaTryLocalFirst  = true
  val faunaContainerImage = "fauna/faunadb:latest"

  private[this] val faunaExposedPort        = 8443
  private[this] val faunaGraphQLExposedPort = 8084
  private[this] lazy val faunaDB: Option[FaunaDB] =
    if (faunaTryLocalFirst) {
      val sock = new Socket()
      Try(sock.connect(new InetSocketAddress("127.0.0.1", faunaExposedPort))) match {
        case Success(_) =>
          sock.close()
          None
        case Failure(_) =>
          startFauna
      }
    } else startFauna

  private def startFauna: Option[FaunaDB] = Some(
    new FaunaDB(
      imageName = faunaContainerImage,
      faunaExposedPort = faunaExposedPort,
      faunaGraphQLExposedPort = faunaGraphQLExposedPort
    )
  )

  lazy val faunaIP: String              = faunaDB.map(_.ipAddress).getOrElse("127.0.0.1")
  lazy val faunaPort: Int               = faunaDB.map(_.mappedFaunaPort).getOrElse(faunaExposedPort)
  lazy val faunaGraphQLPort: Int        = faunaDB.map(_.mappedFaunaGraphQLPort).getOrElse(faunaGraphQLExposedPort)
  lazy val container: Container         = faunaDB.map(_.container).getOrElse(EmptyContainer)
  lazy val faunaEndpoint: String        = s"http://$faunaIP:$faunaPort"
  lazy val faunaGraphQLEndpoint: String = s"http://$faunaIP:$faunaGraphQLPort"
  lazy val faunaSecret: String          = "secret"
  lazy val faunaClient: FaunaClient     = FaunaClient.builder().withEndpoint(faunaEndpoint).withSecret(faunaSecret).build()

  final class FaunaDB(imageName: String, faunaExposedPort: Int, faunaGraphQLExposedPort: Int) {
    val logger      = LoggerFactory.getLogger("faunadb")
    val logConsumer = new Slf4jLogConsumer(logger)
    lazy val container: GenericContainer = GenericContainer(
      dockerImage = imageName,
      exposedPorts = Seq(faunaExposedPort, faunaGraphQLExposedPort),
      waitStrategy = {
        new LoggingLogMessageWaitStrategy(".*FaunaDB is ready.*\\s", 2, Some(logConsumer))
          .withStartupTimeout(Duration.of(60, SECONDS))
      }
    )
    lazy val ipAddress              = container.containerIpAddress
    lazy val mappedFaunaPort        = container.mappedPort(faunaExposedPort)
    lazy val mappedFaunaGraphQLPort = container.mappedPort(faunaGraphQLExposedPort)
  }

  object EmptyContainer extends Container {
    def finished()(implicit description: Description): Unit           = {}
    def failed(e: Throwable)(implicit description: Description): Unit = {}
    def starting()(implicit description: Description): Unit           = {}
    def succeeded()(implicit description: Description): Unit          = {}
  }

  final class LoggingLogMessageWaitStrategy(regEx: String, times: Int, consumer: Option[Consumer[OutputFrame]] = None)
      extends AbstractWaitStrategy {

    override def waitUntilReady(): Unit = {
      val waitingConsumer = new WaitingConsumer
      LogUtils.followOutput(DockerClientFactory.instance.client, waitStrategyTarget.getContainerId, waitingConsumer)

      val waitPredicate = new Predicate[OutputFrame] {
        override def test(t: OutputFrame): Boolean = {
          consumer.foreach(_.accept(t))
          // (?s) enables line terminator matching (equivalent to Pattern.DOTALL)
          t.getUtf8String.matches("(?s)" + regEx)
        }
      }

      try waitingConsumer.waitUntil(waitPredicate, startupTimeout.getSeconds, TimeUnit.SECONDS, times)
      catch {
        case e: TimeoutException =>
          throw new ContainerLaunchException("Timed out waiting for log output matching '" + regEx + "'")
      }
    }

  }

}

MySpec.scala

import org.specs2.mutable.Specification
import org.specs2.specification.BeforeAfterAll

import org.junit.runner.Description

class MySpec extends Specification with BeforeAfterAll with FaunaDBTesting {
  implicit private val suiteDescription = Description.createSuiteDescription(this.getClass)

  override def beforeAll(): Unit = container.starting()
  override def afterAll(): Unit  = container.finished()

  "My Spec" should {
    "query fauna" in faunaClient.query(???)
  }
}