Flyway ZIO Test Aspect

FlywayTestAspect.scala

import java.util.function.UnaryOperator
import javax.sql.DataSource

import org.flywaydb.core.Flyway
import org.flywaydb.core.api.configuration.FluentConfiguration
import org.flywaydb.core.api.output.MigrateResult
import zio.*
import zio.test.TestAspect

object FlywayTestAspect:

  val migrateBefore: TestAspect[Nothing, DataSource, Nothing, Any] =
    migrateBefore()

  def migrateBefore(locations: String*): TestAspect[Nothing, DataSource, Nothing, Any] =
    TestAspect.before(migrate(locations))

  def migrateBefore(configure: UnaryOperator[FluentConfiguration]): TestAspect[Nothing, DataSource, Nothing, Any] =
    TestAspect.before(migrate(configure))

  val migrateBeforeAll: TestAspect[Nothing, DataSource, Nothing, Any] =
    migrateBeforeAll()

  def migrateBeforeAll(locations: String*): TestAspect[Nothing, DataSource, Nothing, Any] =
    TestAspect.beforeAll(migrate(locations))

  def migrateBeforeAll(configure: UnaryOperator[FluentConfiguration]): TestAspect[Nothing, DataSource, Nothing, Any] =
    TestAspect.beforeAll(migrate(configure))

  private def migrate(locations: Seq[String]): URIO[DataSource, MigrateResult] =
    migrate: flywayConfiguration =>
      if locations.isEmpty
      then flywayConfiguration
      else flywayConfiguration.locations(locations*)

  private def migrate(configure: UnaryOperator[FluentConfiguration]): URIO[DataSource, MigrateResult] =
    ZIO
      .serviceWithZIO[DataSource]: dataSource =>
        ZIO.attemptBlockingIO:
          val baseConfiguration: FluentConfiguration =
            Flyway.configure
              .dataSource(dataSource)
              .loggers("slf4j")

          configure(baseConfiguration).load.migrate
      .orDie