TestContainerForAll

apache · Dec 15, 2023 · aa63e26 · aa63e26
1 parent 9047151
commit aa63e26
Show file tree

Hide file tree

Showing 7 changed files with 36 additions and 52 deletions.
diff --git a/externals/kyuubi-jdbc-engine/pom.xml b/externals/kyuubi-jdbc-engine/pom.xml
@@ -64,6 +64,12 @@
             <scope>test</scope>
         </dependency>
 
+        <dependency>
+            <groupId>com.dimafeng</groupId>
+            <artifactId>testcontainers-scala-postgresql_${scala.binary.version}</artifactId>
+            <scope>test</scope>
+        </dependency>
+
         <dependency>
             <groupId>org.apache.kyuubi</groupId>
             <artifactId>${hive.jdbc.artifact}</artifactId>

diff --git a/...bi-jdbc-engine/src/test/scala/org/apache/kyuubi/engine/jdbc/WithJdbcServerContainer.scala b/...bi-jdbc-engine/src/test/scala/org/apache/kyuubi/engine/jdbc/WithJdbcServerContainer.scala
@@ -16,8 +16,8 @@
  */
 package org.apache.kyuubi.engine.jdbc
 
-import com.dimafeng.testcontainers.ForAllTestContainer
+import com.dimafeng.testcontainers.scalatest.TestContainerForAll
 
 import org.apache.kyuubi.KyuubiFunSuite
 
-trait WithJdbcServerContainer extends KyuubiFunSuite with ForAllTestContainer {}
+trait WithJdbcServerContainer extends KyuubiFunSuite with TestContainerForAll {}
diff --git a/...i-jdbc-engine/src/test/scala/org/apache/kyuubi/engine/jdbc/doris/WithDorisContainer.scala b/...i-jdbc-engine/src/test/scala/org/apache/kyuubi/engine/jdbc/doris/WithDorisContainer.scala
@@ -35,7 +35,7 @@ trait WithDorisContainer extends WithJdbcServerContainer {
 
   private val DORIS_BE_SERVICE_NAME = "doris-be"
 
-  override val container: DockerComposeContainer =
+  override val containerDef: DockerComposeContainer.Def =
     DockerComposeContainer
       .Def(
         composeFiles = new File(Utils.getContextOrKyuubiClassLoader
@@ -51,17 +51,11 @@ trait WithDorisContainer extends WithJdbcServerContainer {
             DORIS_BE_PORT,
             waitStrategy =
               new DockerHealthcheckWaitStrategy().withStartupTimeout(Duration.ofMinutes(5)))))
-      .createContainer()
 
-  protected def feUrl: String = {
+  protected def feUrl: String = withContainers { container =>
     val feHost: String = container.getServiceHost(DORIS_FE_SERVICE_NAME, DORIS_FE_PORT)
     val fePort: Int = container.getServicePort(DORIS_FE_SERVICE_NAME, DORIS_FE_PORT)
     val url = s"$feHost:$fePort"
     url
   }
-
-  override def afterAll(): Unit = {
-    super.afterAll()
-    container.close()
-  }
 }
diff --git a/...bc-engine/src/test/scala/org/apache/kyuubi/engine/jdbc/phoenix/WithPhoenixContainer.scala b/...bc-engine/src/test/scala/org/apache/kyuubi/engine/jdbc/phoenix/WithPhoenixContainer.scala
@@ -16,7 +16,7 @@
  */
 package org.apache.kyuubi.engine.jdbc.phoenix
 
-import com.dimafeng.testcontainers.{GenericContainer, SingleContainer}
+import com.dimafeng.testcontainers.GenericContainer
 import org.testcontainers.containers.wait.strategy.Wait
 
 import org.apache.kyuubi.engine.jdbc.WithJdbcServerContainer
@@ -27,21 +27,15 @@ trait WithPhoenixContainer extends WithJdbcServerContainer {
 
   private val phoenixDockerImage = "iteblog/hbase-phoenix-docker:1.0"
 
-  override val container: SingleContainer[_] = GenericContainer(
+  override val containerDef = GenericContainer.Def(
     dockerImage = phoenixDockerImage,
     exposedPorts = Seq(PHOENIX_PORT),
     waitStrategy = Wait.forListeningPort)
 
-  protected def queryServerUrl: String = {
+  protected def queryServerUrl: String = withContainers { container =>
     val queryServerHost: String = container.host
     val queryServerPort: Int = container.mappedPort(PHOENIX_PORT)
     val url = s"$queryServerHost:$queryServerPort"
     url
   }
-
-  override def afterAll(): Unit = {
-    super.afterAll()
-    container.close()
-  }
-
 }
diff --git a/...ine/src/test/scala/org/apache/kyuubi/engine/jdbc/postgresql/WithPostgreSQLContainer.scala b/...ine/src/test/scala/org/apache/kyuubi/engine/jdbc/postgresql/WithPostgreSQLContainer.scala
@@ -16,34 +16,16 @@
  */
 package org.apache.kyuubi.engine.jdbc.postgresql
 
-import com.dimafeng.testcontainers.{GenericContainer, SingleContainer}
-import org.testcontainers.containers.wait.strategy.Wait
+import com.dimafeng.testcontainers.PostgreSQLContainer
+import org.testcontainers.utility.DockerImageName
 
 import org.apache.kyuubi.engine.jdbc.WithJdbcServerContainer
 
 trait WithPostgreSQLContainer extends WithJdbcServerContainer {
 
-  private val POSTGRESQL_PORT = 5432
-
-  private val postgreSQLDockerImage = "postgres"
-
-  override val container: SingleContainer[_] = GenericContainer(
-    dockerImage = postgreSQLDockerImage,
-    exposedPorts = Seq(POSTGRESQL_PORT),
-    env = Map[String, String](
-      "POSTGRES_PASSWORD" -> "postgres"),
-    waitStrategy = Wait.forListeningPort)
-
-  protected def queryUrl: String = {
-    val queryServerHost: String = container.host
-    val queryServerPort: Int = container.mappedPort(POSTGRESQL_PORT)
-    val url = s"$queryServerHost:$queryServerPort"
-    url
-  }
-
-  override def afterAll(): Unit = {
-    super.afterAll()
-    container.close()
-  }
-
+  val containerDef = PostgreSQLContainer.Def(
+    dockerImageName = DockerImageName.parse(s"postgres:16"),
+    databaseName = "postgres",
+    username = "kyuubi",
+    password = "postgres")
 }
diff --git a/...engine/src/test/scala/org/apache/kyuubi/engine/jdbc/postgresql/WithPostgreSQLEngine.scala b/...engine/src/test/scala/org/apache/kyuubi/engine/jdbc/postgresql/WithPostgreSQLEngine.scala
@@ -21,13 +21,15 @@ import org.apache.kyuubi.engine.jdbc.WithJdbcEngine
 
 trait WithPostgreSQLEngine extends WithJdbcEngine with WithPostgreSQLContainer {
 
-  override def withKyuubiConf: Map[String, String] = Map(
-    ENGINE_SHARE_LEVEL.key -> "SERVER",
-    ENGINE_JDBC_CONNECTION_URL.key -> s"jdbc:postgresql://$queryUrl/postgres",
-    ENGINE_JDBC_CONNECTION_USER.key -> "postgres",
-    ENGINE_JDBC_CONNECTION_PASSWORD.key -> "postgres",
-    ENGINE_TYPE.key -> "jdbc",
-    ENGINE_JDBC_SHORT_NAME.key -> "postgresql",
-    ENGINE_JDBC_DRIVER_CLASS.key -> "org.postgresql.Driver")
+  override def withKyuubiConf: Map[String, String] = withContainers { container =>
+    Map(
+      ENGINE_SHARE_LEVEL.key -> "SERVER",
+      ENGINE_JDBC_CONNECTION_URL.key -> container.jdbcUrl,
+      ENGINE_JDBC_CONNECTION_USER.key -> container.username,
+      ENGINE_JDBC_CONNECTION_PASSWORD.key -> container.password,
+      ENGINE_TYPE.key -> "jdbc",
+      ENGINE_JDBC_SHORT_NAME.key -> "postgresql",
+      ENGINE_JDBC_DRIVER_CLASS.key -> container.driverClassName)
+  }
 
 }
diff --git a/pom.xml b/pom.xml
@@ -574,6 +574,12 @@
                 <version>${testcontainers-scala.version}</version>
             </dependency>
 
+            <dependency>
+                <groupId>com.dimafeng</groupId>
+                <artifactId>testcontainers-scala-postgresql_${scala.binary.version}</artifactId>
+                <version>${testcontainers-scala.version}</version>
+            </dependency>
+
             <dependency>
                 <groupId>com.dimafeng</groupId>
                 <artifactId>testcontainers-scala-trino_${scala.binary.version}</artifactId>