Try to fix unit test

linghengqian · Nov 30, 2024 · f3041ad · f3041ad
1 parent 8cecb18
commit f3041ad
Show file tree

Hide file tree

Showing 2 changed files with 24 additions and 26 deletions.
diff --git a/...thin/src/test/java/io/github/linghengqian/hive/server2/jdbc/driver/thin/HS2Container.java b/...thin/src/test/java/io/github/linghengqian/hive/server2/jdbc/driver/thin/HS2Container.java
@@ -7,8 +7,10 @@
 import java.io.IOException;
 import java.net.ServerSocket;
 
-@SuppressWarnings({"resource", "OctalInteger"})
+@SuppressWarnings({"resource"})
 public class HS2Container extends GenericContainer<HS2Container> {
+
+    String zookeeperConnectionString;
     private static final String STARTER_SCRIPT = "/testcontainers_start.sh";
     private final int randomPortFirst = getRandomPort();
 
@@ -21,14 +23,18 @@ public HS2Container(final String dockerImageName) {
         );
     }
 
+    public HS2Container withZookeeperConnectionString(final String zookeeperConnectionString) {
+        this.zookeeperConnectionString = zookeeperConnectionString;
+        return self();
+    }
+
     @Override
     protected void containerIsStarting(InspectContainerResponse containerInfo) {
-        Integer mappedPort = getMappedPort(randomPortFirst);
         String command = """
                 #!/bin/bash
-                export SERVICE_OPTS='-Dhive.server2.support.dynamic.service.discovery=true -Dhive.zookeeper.quorum=foo:2181 -Dhive.server2.thrift.bind.host=0.0.0.0 -Dhive.server2.thrift.port=%s'
-                /usr/local/bin/docker-entrypoint.sh
-                """.formatted(mappedPort);
+                export SERVICE_OPTS="-Dhive.server2.support.dynamic.service.discovery=true -Dhive.zookeeper.quorum=%s -Dhive.server2.thrift.bind.host=0.0.0.0 -Dhive.server2.thrift.port=%s"
+                /entrypoint.sh
+                """.formatted(zookeeperConnectionString, getMappedPort(randomPortFirst));
         copyFileToContainer(Transferable.of(command, 0777), STARTER_SCRIPT);
     }
 

diff --git a/...a/io/github/linghengqian/hive/server2/jdbc/driver/thin/ZookeeperServiceDiscoveryTest.java b/...a/io/github/linghengqian/hive/server2/jdbc/driver/thin/ZookeeperServiceDiscoveryTest.java
@@ -61,21 +61,18 @@ static void afterAll() {
 
     @Test
     void assertShardingInLocalTransactions() throws SQLException {
-        GenericContainer<?> hs2FirstContainer = new HS2Container("apache/hive:4.0.1")
+        try (GenericContainer<?> hs2FirstContainer = new HS2Container("apache/hive:4.0.1")
                 .withNetwork(NETWORK)
-                .dependsOn(ZOOKEEPER_CONTAINER);
-        hs2FirstContainer.start();
-        awaitHS2(hs2FirstContainer.getFirstMappedPort());
-        DataSource dataSource = createDataSource();
-        extractedSQL(dataSource);
-        hs2FirstContainer.stop();
-    }
-
-    private DataSource createDataSource() {
-        HikariConfig config = new HikariConfig();
-        config.setDriverClassName("org.apache.hive.jdbc.HiveDriver");
-        config.setJdbcUrl(jdbcUrlPrefix + jdbcUrlSuffix);
-        return new HikariDataSource(config);
+                .withZookeeperConnectionString(ZOOKEEPER_CONTAINER.getNetworkAliases().get(0) + ":" + ZOOKEEPER_CONTAINER.getMappedPort(2181))
+                .dependsOn(ZOOKEEPER_CONTAINER)) {
+            hs2FirstContainer.start();
+            awaitHS2(hs2FirstContainer.getFirstMappedPort());
+            HikariConfig config = new HikariConfig();
+            config.setDriverClassName("org.apache.hive.jdbc.HiveDriver");
+            config.setJdbcUrl(jdbcUrlPrefix + jdbcUrlSuffix);
+            DataSource dataSource = new HikariDataSource(config);
+            extractedSQL(dataSource);
+        }
     }
 
     private static void extractedSQL(final DataSource dataSource) throws SQLException {
@@ -96,14 +93,9 @@ private static void extractedSQL(final DataSource dataSource) throws SQLExceptio
         }
     }
 
-    private Connection openConnection() throws SQLException {
-        Properties props = new Properties();
-        return DriverManager.getConnection(jdbcUrlPrefix + jdbcUrlSuffix, props);
-    }
-
     private void awaitHS2(final int hiveServer2Port) {
         String connectionString = ZOOKEEPER_CONTAINER.getHost() + ":" + ZOOKEEPER_CONTAINER.getMappedPort(2181);
-        await().atMost(Duration.ofMinutes(2L)).ignoreExceptions().until(() -> {
+        await().atMost(Duration.ofMinutes(1L)).ignoreExceptions().until(() -> {
             try (CuratorFramework client = CuratorFrameworkFactory.builder()
                     .connectString(connectionString)
                     .retryPolicy(new ExponentialBackoffRetry(1000, 3))
@@ -115,7 +107,7 @@ private void awaitHS2(final int hiveServer2Port) {
             }
         });
         await().atMost(Duration.ofMinutes(1L)).ignoreExceptions().until(() -> {
-            openConnection().close();
+            DriverManager.getConnection(jdbcUrlPrefix + jdbcUrlSuffix, new Properties()).close();
             return true;
         });
     }