chore: refactor AuroraTestUtility

wrapper/src/test/java/integration/container/tests/AutoscalingTests.java

@@ -48,6 +48,7 @@
 import org.junit.jupiter.api.TestTemplate;
 import org.junit.jupiter.api.extension.ExtendWith;
 import software.amazon.jdbc.ConnectionProviderManager;
+import software.amazon.jdbc.Driver;
 import software.amazon.jdbc.HikariPoolConfigurator;
 import software.amazon.jdbc.HikariPooledConnectionProvider;
 import software.amazon.jdbc.PropertyDefinition;
@@ -116,7 +117,7 @@ public void test_pooledConnectionAutoScaling_setReadOnlyOnOldConnection()
             null,
             poolExpirationNanos,
             TimeUnit.MINUTES.toNanos(10));
-    ConnectionProviderManager.setConnectionProvider(provider);
+    Driver.setCustomConnectionProvider(provider);
 
     final List<Connection> connections = new ArrayList<>();
     try {
@@ -127,8 +128,11 @@ public void test_pooledConnectionAutoScaling_setReadOnlyOnOldConnection()
       }
 
       final Connection newInstanceConn;
+      final String instanceClass =
+          auroraUtil.getDbInstanceClass(TestEnvironment.getCurrent().getInfo().getRequest());
       final TestInstanceInfo newInstance =
-          auroraUtil.createInstance("auto-scaling-instance");
+          auroraUtil.createInstance(instanceClass, "auto-scaling-instance");
+      instances.add(newInstance);
       try {
         newInstanceConn =
             DriverManager.getConnection(ConnectionStringHelper.getWrapperUrl(), props);
@@ -144,6 +148,7 @@ public void test_pooledConnectionAutoScaling_setReadOnlyOnOldConnection()
         newInstanceConn.setReadOnly(false);
       } finally {
         auroraUtil.deleteInstance(newInstance);
+        instances.remove(newInstance);
       }
 
       final long deletionCheckTimeout = System.nanoTime() + TimeUnit.MINUTES.toNanos(5);
@@ -200,8 +205,11 @@ public void test_pooledConnectionAutoScaling_failoverFromDeletedReader()
       }
 
       final Connection newInstanceConn;
+      final String instanceClass =
+          auroraUtil.getDbInstanceClass(TestEnvironment.getCurrent().getInfo().getRequest());
       final TestInstanceInfo newInstance =
-          auroraUtil.createInstance("auto-scaling-instance");
+          auroraUtil.createInstance(instanceClass, "auto-scaling-instance");
+      instances.add(newInstance);
       try {
         newInstanceConn =
             DriverManager.getConnection(ConnectionStringHelper.getWrapperUrl(newInstance), props);
@@ -214,6 +222,7 @@ public void test_pooledConnectionAutoScaling_failoverFromDeletedReader()
             .anyMatch((url) -> url.equals(newInstance.getUrl())));
       } finally {
         auroraUtil.deleteInstance(newInstance);
+        instances.remove(newInstance);
       }
 
       auroraUtil.assertFirstQueryThrows(newInstanceConn, FailoverSuccessSQLException.class);

wrapper/src/test/java/integration/container/tests/AwsIamIntegrationTest.java

@@ -40,7 +40,6 @@
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.MethodOrderer;
 import org.junit.jupiter.api.Order;
-import org.junit.jupiter.api.Test;
 import org.junit.jupiter.api.TestMethodOrder;
 import org.junit.jupiter.api.TestTemplate;
 import org.junit.jupiter.api.extension.ExtendWith;

wrapper/src/test/java/integration/container/tests/HikariTests.java

@@ -90,6 +90,7 @@
 public class HikariTests {
 
   private static final Logger LOGGER = Logger.getLogger(HikariTests.class.getName());
+  protected static final AuroraTestUtility auroraUtil = AuroraTestUtility.getUtility();
 
   @TestTemplate
   public void testOpenConnectionWithUrl() throws SQLException {
@@ -176,8 +177,6 @@ public void testOpenConnectionWithDataSourceClassName() throws SQLException {
   @EnableOnTestFeature({TestEnvironmentFeatures.NETWORK_OUTAGES_ENABLED})
   @EnableOnNumOfInstances(min = 3)
   public void testFailoverLostConnection() throws SQLException {
-
-    final AuroraTestUtility auroraUtil = AuroraTestUtility.getUtility();
     final Properties customProps = new Properties();
     PLUGINS.set(customProps, "failover");
     FAILOVER_TIMEOUT_MS.set(customProps, Integer.toString(1));
@@ -218,8 +217,6 @@ public void testFailoverLostConnection() throws SQLException {
   @EnableOnTestFeature({TestEnvironmentFeatures.NETWORK_OUTAGES_ENABLED})
   @EnableOnNumOfInstances(min = 3)
   public void testEFMFailover() throws SQLException {
-
-    final AuroraTestUtility auroraUtil = AuroraTestUtility.getUtility();
     ProxyHelper.disableAllConnectivity();
 
     final List<TestInstanceInfo> instances = TestEnvironment.getCurrent()
@@ -270,7 +267,6 @@ public void testEFMFailover() throws SQLException {
   @EnableOnNumOfInstances(min = 2)
   public void testInternalPools_driverWriterFailoverOnGetConnectionInvocation()
       throws SQLException, InterruptedException {
-    final AuroraTestUtility auroraUtil = AuroraTestUtility.getUtility();
     final TestProxyDatabaseInfo proxyInfo = TestEnvironment.getCurrent().getInfo().getProxyDatabaseInfo();
     final List<TestInstanceInfo> instances = proxyInfo.getInstances();
     final TestInstanceInfo reader = instances.get(1);
@@ -311,7 +307,6 @@ public void testInternalPools_driverWriterFailoverOnGetConnectionInvocation()
   @EnableOnNumOfInstances(min = 2)
   public void testInternalPools_driverReaderFailoverOnGetConnectionInvocation()
       throws SQLException, InterruptedException {
-    final AuroraTestUtility auroraUtil = AuroraTestUtility.getUtility();
     final TestProxyDatabaseInfo proxyInfo = TestEnvironment.getCurrent().getInfo().getProxyDatabaseInfo();
     final List<TestInstanceInfo> instances = proxyInfo.getInstances();
     final TestInstanceInfo writer = instances.get(0);
@@ -360,7 +355,6 @@ public void testInternalPools_driverReaderFailoverOnGetConnectionInvocation()
   @EnableOnNumOfInstances(max = 1)
   public void testInternalPools_driverWriterFailoverOnGetConnectionInvocation_singleInstance()
       throws SQLException, InterruptedException {
-    final AuroraTestUtility auroraUtil = AuroraTestUtility.getUtility();
     final TestProxyDatabaseInfo proxyInfo = TestEnvironment.getCurrent().getInfo().getProxyDatabaseInfo();
     final List<TestInstanceInfo> instances = proxyInfo.getInstances();
     final TestInstanceInfo writer = instances.get(0);

wrapper/src/test/java/integration/container/tests/ReadWriteSplittingPerformanceTest.java

@@ -17,7 +17,6 @@
 package integration.container.tests;
 
 import com.zaxxer.hikari.HikariConfig;
-import integration.DriverHelper;
 import integration.TestEnvironmentFeatures;
 import integration.container.ConnectionStringHelper;
 import integration.container.TestDriverProvider;

wrapper/src/test/java/integration/container/tests/ReadWriteSplittingTests.java

@@ -569,8 +569,8 @@ public void test_failoverReaderToWriter_setReadOnlyTrueFalse()
 
       ProxyHelper.enableAllConnectivity();
 
-      // During failover above some of readers have been tried to connect to and failed since they were not
-      // available. We should expect that some of readers in topology are marked as UNAVAILABLE.
+      // During failover above some of the readers have been tried to connect to and failed since they were not
+      // available. We should expect that some of the readers in topology are marked as UNAVAILABLE.
       // The following code reset node availability and make them AVAILABLE.
       // That is important for further steps.
       TestPluginServiceImpl.clearHostAvailabilityCache();
@@ -857,7 +857,7 @@ public void test_pooledConnection_differentUsers() throws SQLException {
 
       assertThrows(
           HikariPool.PoolInitializationException.class, () -> {
-            try (final Connection conn = DriverManager.getConnection(
+            try (final Connection ignored = DriverManager.getConnection(
                 ConnectionStringHelper.getWrapperUrl(), wrongUserRightPasswordProps)) {
               // Do nothing (close connection automatically)
             }

wrapper/src/test/java/integration/host/TestEnvironment.java

@@ -55,6 +55,7 @@
 import org.testcontainers.shaded.org.apache.commons.lang3.NotImplementedException;
 import org.testcontainers.utility.MountableFile;
 import software.amazon.awssdk.services.rds.model.DBCluster;
+import software.amazon.awssdk.services.rds.model.DBInstance;
 import software.amazon.jdbc.util.StringUtils;
 
 public class TestEnvironment implements AutoCloseable {
@@ -174,8 +175,7 @@ public static TestEnvironment build(TestEnvironmentRequest request) throws IOExc
     return env;
   }
 
-  private static TestEnvironment createAuroraOrMultiAzEnvironment(TestEnvironmentRequest request)
-      throws URISyntaxException {
+  private static TestEnvironment createAuroraOrMultiAzEnvironment(TestEnvironmentRequest request) {
 
     EnvPreCreateInfo preCreateInfo =
         TestEnvironmentProvider.preCreateInfos.get(request.getEnvPreCreateIndex());
@@ -325,7 +325,7 @@ private static void createDatabaseContainers(TestEnvironment env) {
     }
   }
 
-  private static void createDbCluster(TestEnvironment env) throws URISyntaxException {
+  private static void createDbCluster(TestEnvironment env) {
 
     switch (env.info.getRequest().getDatabaseInstances()) {
       case SINGLE_INSTANCE:
@@ -351,7 +351,7 @@ private static void createDbCluster(TestEnvironment env) throws URISyntaxExcepti
     }
   }
 
-  private static void createDbCluster(TestEnvironment env, int numOfInstances) throws URISyntaxException {
+  private static void createDbCluster(TestEnvironment env, int numOfInstances) {
 
     env.info.setRegion(
         !StringUtils.isNullOrEmpty(config.rdsDbRegion)
@@ -371,9 +371,6 @@ private static void createDbCluster(TestEnvironment env, int numOfInstances) thr
             env.awsAccessKeyId,
             env.awsSecretAccessKey,
             env.awsSessionToken);
-
-    ArrayList<TestInstanceInfo> instances = new ArrayList<>();
-
     if (env.reuseAuroraDbCluster) {
       if (StringUtils.isNullOrEmpty(env.auroraClusterDomain)) {
         throw new RuntimeException("Environment variable RDS_CLUSTER_DOMAIN is required.");
@@ -403,8 +400,6 @@ private static void createDbCluster(TestEnvironment env, int numOfInstances) thr
 
       env.info.setDatabaseEngine(clusterInfo.engine());
       env.info.setDatabaseEngineVersion(clusterInfo.engineVersion());
-      instances.addAll(env.auroraUtil.getClusterInstanceIds(env.auroraClusterName));
-
     } else {
       if (StringUtils.isNullOrEmpty(env.auroraClusterName)) {
         int remainingTries = 5;
@@ -431,23 +426,30 @@ private static void createDbCluster(TestEnvironment env, int numOfInstances) thr
         if (StringUtils.isNullOrEmpty(engineVersion)) {
           throw new RuntimeException("Failed to get engine version.");
         }
-        String instanceClass = getDbInstanceClass(env.info.getRequest());
+        String instanceClass = env.auroraUtil.getDbInstanceClass(env.info.getRequest());
 
         LOGGER.finer(
             "Using " + engine + " " + engineVersion);
 
-        env.auroraClusterDomain =
-            env.auroraUtil.createCluster(
-                env.info.getDatabaseInfo().getUsername(),
-                env.info.getDatabaseInfo().getPassword(),
-                env.info.getDatabaseInfo().getDefaultDbName(),
-                env.auroraClusterName,
-                env.info.getRequest().getDatabaseEngineDeployment(),
-                engine,
-                instanceClass,
-                engineVersion,
-                numOfInstances,
-                instances);
+        env.auroraUtil.createCluster(
+            env.info.getDatabaseInfo().getUsername(),
+            env.info.getDatabaseInfo().getPassword(),
+            env.info.getDatabaseInfo().getDefaultDbName(),
+            env.auroraClusterName,
+            env.info.getRequest().getDatabaseEngineDeployment(),
+            env.info.getRegion(),
+            engine,
+            instanceClass,
+            engineVersion,
+            numOfInstances);
+
+        List<DBInstance> dbInstances = env.auroraUtil.getDBInstances(env.auroraClusterName);
+        if (dbInstances.isEmpty()) {
+          throw new RuntimeException("Failed to get instance information for cluster " + env.auroraClusterName);
+        }
+
+        final String instanceEndpoint = dbInstances.get(0).endpoint().address();
+        env.auroraClusterDomain = instanceEndpoint.substring(instanceEndpoint.indexOf(".") + 1);
         env.info.setDatabaseEngine(engine);
         env.info.setDatabaseEngineVersion(engineVersion);
         LOGGER.finer(
@@ -458,7 +460,7 @@ private static void createDbCluster(TestEnvironment env, int numOfInstances) thr
 
         // remove cluster and instances
         LOGGER.finer("Deleting cluster " + env.auroraClusterName);
-        env.auroraUtil.deleteCluster(env.auroraClusterName);
+        env.auroraUtil.deleteCluster(env.auroraClusterName, env.info.getRequest().getDatabaseEngineDeployment());
         LOGGER.finer("Deleted cluster " + env.auroraClusterName);
 
         throw new RuntimeException(e);
@@ -478,6 +480,7 @@ private static void createDbCluster(TestEnvironment env, int numOfInstances) thr
             env.auroraClusterName + ".cluster-ro-" + env.auroraClusterDomain, port);
     env.info.getDatabaseInfo().setInstanceEndpointSuffix(env.auroraClusterDomain, port);
 
+    List<TestInstanceInfo> instances = env.auroraUtil.getTestInstancesInfo(env.auroraClusterName);
     env.info.getDatabaseInfo().getInstances().clear();
     env.info.getDatabaseInfo().getInstances().addAll(instances);
 
@@ -586,18 +589,6 @@ private static String findEngineVersion(
     }
   }
 
-  private static String getDbInstanceClass(TestEnvironmentRequest request) {
-    switch (request.getDatabaseEngineDeployment()) {
-      case AURORA:
-        return "db.r5.large";
-      case RDS:
-      case RDS_MULTI_AZ_CLUSTER:
-        return "db.m5d.large";
-      default:
-        throw new NotImplementedException(request.getDatabaseEngine().toString());
-    }
-  }
-
   private static int getPort(TestEnvironmentRequest request) {
     switch (request.getDatabaseEngine()) {
       case MYSQL:
@@ -979,16 +970,14 @@ private String[] buildHibernateCommands(boolean debugMode) {
 
   @Override
   public void close() throws Exception {
-    if (this.databaseContainers != null) {
-      for (GenericContainer<?> container : this.databaseContainers) {
-        try {
-          container.stop();
-        } catch (Exception ex) {
-          // ignore
-        }
+    for (GenericContainer<?> container : this.databaseContainers) {
+      try {
+        container.stop();
+      } catch (Exception ex) {
+        // ignore
       }
-      this.databaseContainers.clear();
     }
+    this.databaseContainers.clear();
 
     if (this.telemetryXRayContainer != null) {
       this.telemetryXRayContainer.stop();
@@ -1035,7 +1024,7 @@ private void deleteDbCluster() {
 
     if (!this.reuseAuroraDbCluster) {
       LOGGER.finest("Deleting cluster " + this.auroraClusterName + ".cluster-" + this.auroraClusterDomain);
-      auroraUtil.deleteCluster(this.auroraClusterName);
+      auroraUtil.deleteCluster(this.auroraClusterName, this.info.getRequest().getDatabaseEngineDeployment());
       LOGGER.finest("Deleted cluster " + this.auroraClusterName + ".cluster-" + this.auroraClusterDomain);
     }
   }