Simplify database migration tests

trinodb · Jan 17, 2025 · 5306731 · 5306731
1 parent 64dcbd5
commit 5306731
Show file tree

Hide file tree

Showing 3 changed files with 31 additions and 62 deletions.
diff --git a/gateway-ha/src/test/java/io/trino/gateway/ha/persistence/BaseTestDatabaseMigrations.java b/gateway-ha/src/test/java/io/trino/gateway/ha/persistence/BaseTestDatabaseMigrations.java
@@ -25,6 +25,8 @@
 
 import java.util.List;
 
+import static java.lang.String.format;
+import static java.util.Objects.requireNonNull;
 import static org.assertj.core.api.Assertions.assertThat;
 import static org.junit.jupiter.api.TestInstance.Lifecycle.PER_CLASS;
 import static org.junit.jupiter.api.parallel.ExecutionMode.SAME_THREAD;
@@ -34,14 +36,19 @@
 @Isolated
 public abstract class BaseTestDatabaseMigrations
 {
-    protected final JdbcDatabaseContainer<?> container = startContainer();
-    protected final Jdbi jdbi = Jdbi.create(container.getJdbcUrl(), container.getUsername(), container.getPassword());
-
-    protected abstract JdbcDatabaseContainer<?> startContainer();
+    protected abstract void createGatewaySchema();
 
-    protected abstract String getDriver();
+    private final JdbcDatabaseContainer<?> container;
+    private final String schema;
+    protected final Jdbi jdbi;
 
-    protected abstract void createGatewaySchema();
+    public BaseTestDatabaseMigrations(JdbcDatabaseContainer<?> container, String schema)
+    {
+        this.container = requireNonNull(container, "container is null");
+        this.container.start();
+        this.schema = requireNonNull(schema, "schema is null");
+        jdbi = Jdbi.create(container.getJdbcUrl(), container.getUsername(), container.getPassword());
+    }
 
     @AfterAll
     public final void close()
@@ -52,13 +59,7 @@ public final void close()
     @Test
     public void testMigrationWithEmptyDatabase()
     {
-        DataStoreConfiguration config = new DataStoreConfiguration(
-                container.getJdbcUrl(),
-                container.getUsername(),
-                container.getPassword(),
-                getDriver(),
-                4,
-                true);
+        DataStoreConfiguration config = dataStoreConfiguration();
         FlywayMigration.migrate(config);
         verifyGatewaySchema(0);
 
@@ -68,13 +69,7 @@ public void testMigrationWithEmptyDatabase()
     @Test
     public void testMigrationWithNonemptyDatabase()
     {
-        DataStoreConfiguration config = new DataStoreConfiguration(
-                container.getJdbcUrl(),
-                container.getUsername(),
-                container.getPassword(),
-                getDriver(),
-                4,
-                true);
+        DataStoreConfiguration config = dataStoreConfiguration();
         String t1Create = "CREATE TABLE t1 (id INT)";
         String t2Create = "CREATE TABLE t2 (id INT)";
         Handle jdbiHandle = jdbi.open();
@@ -98,13 +93,7 @@ public void testMigrationWithExistingGatewaySchema()
         // add a row to one of the existing tables before migration
         jdbi.withHandle(handle ->
                 handle.execute("INSERT INTO resource_groups_global_properties VALUES ('a_name', 'a_value')"));
-        DataStoreConfiguration config = new DataStoreConfiguration(
-                container.getJdbcUrl(),
-                container.getUsername(),
-                container.getPassword(),
-                getDriver(),
-                4,
-                true);
+        DataStoreConfiguration config = dataStoreConfiguration();
         FlywayMigration.migrate(config);
         verifyGatewaySchema(1);
         dropAllTables();
@@ -137,7 +126,7 @@ protected void dropAllTables()
         String exactMatchTable = "DROP TABLE IF EXISTS exact_match_source_selectors";
         String flywayHistoryTable = "DROP TABLE IF EXISTS flyway_schema_history";
         Handle jdbiHandle = jdbi.open();
-        String sql = String.format("SELECT 1 FROM information_schema.tables WHERE table_schema = '%s'", getTestSchema());
+        String sql = format("SELECT 1 FROM information_schema.tables WHERE table_schema = '%s'", schema);
         verifyResultSetCount(sql, 7);
         jdbiHandle.execute(gatewayBackendTable);
         jdbiHandle.execute(queryHistoryTable);
@@ -150,8 +139,14 @@ protected void dropAllTables()
         jdbiHandle.close();
     }
 
-    protected String getTestSchema()
+    private DataStoreConfiguration dataStoreConfiguration()
     {
-        return "public";
+        return new DataStoreConfiguration(
+                container.getJdbcUrl(),
+                container.getUsername(),
+                container.getPassword(),
+                container.getDriverClassName(),
+                4,
+                true);
     }
 }
diff --git a/gateway-ha/src/test/java/io/trino/gateway/ha/persistence/TestDatabaseMigrationsMySql.java b/gateway-ha/src/test/java/io/trino/gateway/ha/persistence/TestDatabaseMigrationsMySql.java
@@ -14,30 +14,14 @@
 package io.trino.gateway.ha.persistence;
 
 import org.jdbi.v3.core.Handle;
-import org.testcontainers.containers.JdbcDatabaseContainer;
 import org.testcontainers.containers.MySQLContainer;
 
-public class TestDatabaseMigrationsMySql
+final class TestDatabaseMigrationsMySql
         extends BaseTestDatabaseMigrations
 {
-    @Override
-    protected final JdbcDatabaseContainer<?> startContainer()
-    {
-        JdbcDatabaseContainer<?> container = new MySQLContainer<>("mysql:8.0.36");
-        container.start();
-        return container;
-    }
-
-    @Override
-    protected final String getDriver()
-    {
-        return "com.mysql.cj.jdbc.Driver";
-    }
-
-    @Override
-    protected final String getTestSchema()
+    public TestDatabaseMigrationsMySql()
     {
-        return "test";
+        super(new MySQLContainer<>("mysql:8.0.36"), "test");
     }
 
     @Override

diff --git a/...nce/TestDatabaseMigrationsPostgresql.java → ...nce/TestDatabaseMigrationsPostgreSql.java b/...nce/TestDatabaseMigrationsPostgresql.java → ...nce/TestDatabaseMigrationsPostgreSql.java
@@ -14,24 +14,14 @@
 package io.trino.gateway.ha.persistence;
 
 import org.jdbi.v3.core.Handle;
-import org.testcontainers.containers.JdbcDatabaseContainer;
 import org.testcontainers.containers.PostgreSQLContainer;
 
-public class TestDatabaseMigrationsPostgresql
+final class TestDatabaseMigrationsPostgreSql
         extends BaseTestDatabaseMigrations
 {
-    @Override
-    protected final JdbcDatabaseContainer<?> startContainer()
-    {
-        JdbcDatabaseContainer<?> container = new PostgreSQLContainer<>("postgres:11");
-        container.start();
-        return container;
-    }
-
-    @Override
-    protected final String getDriver()
+    public TestDatabaseMigrationsPostgreSql()
     {
-        return "org.postgresql.Driver";
+        super(new PostgreSQLContainer<>("postgres:11"), "public");
     }
 
     @Override