From a1d27ebca61b33066b627ba22d77510f1c49aa16 Mon Sep 17 00:00:00 2001
From: Karen Chen <karen.chen@improving.com>
Date: Fri, 15 Mar 2024 22:35:56 -0700
Subject: [PATCH] fix: failing integration tests due to not waiting for TAZ
 cluster to fully failover

---
 .../container/utils/rds_test_utility.py       |  1 +
 .../integration/host/TestEnvironment.java     | 25 ++++++++++++++-----
 2 files changed, 20 insertions(+), 6 deletions(-)

diff --git a/tests/integration/container/utils/rds_test_utility.py b/tests/integration/container/utils/rds_test_utility.py
index 7e38b608..4b440211 100644
--- a/tests/integration/container/utils/rds_test_utility.py
+++ b/tests/integration/container/utils/rds_test_utility.py
@@ -156,6 +156,7 @@ def failover_cluster_and_wait_until_writer_changed(
             sleep(1)
             cluster_address = socket.gethostbyname(cluster_endpoint)
 
+        self.make_sure_instances_up(self.get_instance_ids())
         self.logger.debug("Testing.FinishedFailover", initial_writer_id, str((perf_counter_ns() - start) / 1_000_000))
 
     def failover_cluster(self, cluster_id: Optional[str] = None) -> None:
diff --git a/tests/integration/host/src/test/java/integration/host/TestEnvironment.java b/tests/integration/host/src/test/java/integration/host/TestEnvironment.java
index f99fc3bf..208fbd42 100644
--- a/tests/integration/host/src/test/java/integration/host/TestEnvironment.java
+++ b/tests/integration/host/src/test/java/integration/host/TestEnvironment.java
@@ -16,8 +16,6 @@
 
 package integration.host;
 
-import static org.junit.jupiter.api.Assertions.assertEquals;
-
 import com.fasterxml.jackson.core.JsonProcessingException;
 import com.fasterxml.jackson.databind.ObjectMapper;
 import eu.rekawek.toxiproxy.ToxiproxyClient;
@@ -44,6 +42,7 @@
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.Executors;
 import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicInteger;
 import java.util.logging.Logger;
 import org.testcontainers.containers.GenericContainer;
 import org.testcontainers.containers.Network;
@@ -70,6 +69,7 @@ public class TestEnvironment implements AutoCloseable {
   private static final TestEnvironmentConfiguration config = new TestEnvironmentConfiguration();
   private static final boolean USE_OTLP_CONTAINER_FOR_TRACES = false;
 
+  private static final AtomicInteger ipAddressUsageRefCount = new AtomicInteger(0);
   private final TestEnvironmentInfo info =
       new TestEnvironmentInfo(); // only this info is passed to test container
 
@@ -102,12 +102,22 @@ private TestEnvironment(TestEnvironmentRequest request) {
   }
 
   public static TestEnvironment build(TestEnvironmentRequest request) throws IOException {
+      DatabaseEngineDeployment deployment = request.getDatabaseEngineDeployment();
+    if (deployment == DatabaseEngineDeployment.AURORA
+        || deployment == DatabaseEngineDeployment.RDS
+        || deployment == DatabaseEngineDeployment.RDS_MULTI_AZ) {
+      // These environment require creating external database cluster that should be publicly available.
+      // Corresponding AWS Security Groups should be configured and the test task runner IP address
+      // should be whitelisted.
+      ipAddressUsageRefCount.incrementAndGet();
+    }
+
     LOGGER.finest("Building test env: " + request.getEnvPreCreateIndex());
     preCreateEnvironment(request.getEnvPreCreateIndex());
 
     TestEnvironment env;
 
-    switch (request.getDatabaseEngineDeployment()) {
+    switch (deployment) {
       case DOCKER:
         env = new TestEnvironment(request);
         initDatabaseParams(env);
@@ -185,8 +195,7 @@ private static TestEnvironment createAuroraOrMultiAzEnvironment(TestEnvironmentR
       if (result instanceof Exception) {
         throw new RuntimeException((Exception) result);
       }
-      if (result instanceof TestEnvironment) {
-        TestEnvironment resultTestEnvironment = (TestEnvironment) result;
+      if (result instanceof TestEnvironment resultTestEnvironment) {
         LOGGER.finer(() -> String.format("Use pre-created DB cluster: %s.cluster-%s",
             resultTestEnvironment.auroraClusterName, resultTestEnvironment.auroraClusterDomain));
 
@@ -894,7 +903,11 @@ public void close() throws Exception {
 
   private void deleteDbCluster() {
     if (!this.reuseAuroraDbCluster && !StringUtils.isNullOrEmpty(this.runnerIP)) {
-      auroraUtil.ec2DeauthorizesIP(runnerIP);
+      if (ipAddressUsageRefCount.decrementAndGet() == 0) {
+        // Another test environments are still in use of test task runner IP address.
+        // The last execute tst environment will do the cleanup.
+        auroraUtil.ec2DeauthorizesIP(runnerIP);
+      }
     }
 
     if (!this.reuseAuroraDbCluster) {