fix: stabilize prepared statement invalidation test

Author: dkropachev

integration-tests/src/test/java/com/datastax/oss/driver/core/cql/PreparedStatementCachingIT.java

@@ -18,6 +18,7 @@
 package com.datastax.oss.driver.core.cql;
 
 import static org.assertj.core.api.Assertions.assertThat;
+import static org.awaitility.Awaitility.await;
 
 import com.codahale.metrics.Gauge;
 import com.datastax.oss.driver.api.core.CqlSession;
@@ -226,7 +227,6 @@ private void invalidationTestInner(
       ByteBuffer queryId2 = session.prepare(preparedStmtQueryType2).getId();
       assertThat(getPreparedCacheSize(session)).isEqualTo(2);
 
-      CountDownLatch preparedStmtCacheRemoveLatch = new CountDownLatch(1);
       CountDownLatch typeChangeEventLatch = new CountDownLatch(expectedChangedTypes.size());
 
       DefaultDriverContext ctx = (DefaultDriverContext) session.getContext();
@@ -260,23 +260,24 @@ private void invalidationTestInner(
                   removedQueryEventError.set(
                       Optional.of("Unable to set reference for PS removal event"));
                 }
-                preparedStmtCacheRemoveLatch.countDown();
               });
 
       // alter test_type_2 to trigger cache invalidation and above events
       session.execute("ALTER TYPE test_type_2 add i blob");
 
-      // wait for latches and fail if they don't reach zero before timeout
+      // wait for schema propagation and prepared statement invalidation to settle
       assertThat(Uninterruptibles.awaitUninterruptibly(typeChangeEventLatch, 60, TimeUnit.SECONDS))
           .withFailMessage("typeChangeEventLatch did not trigger before timeout")
           .isTrue();
-      assertThat(
-              Uninterruptibles.awaitUninterruptibly(
-                  preparedStmtCacheRemoveLatch, 60, TimeUnit.SECONDS))
-          .withFailMessage("preparedStmtCacheRemoveLatch did not trigger before timeout")
-          .isTrue();
+      await()
+          .atMost(120, TimeUnit.SECONDS)
+          .untilAsserted(
+              () -> {
+                assertThat(removedQueryIds.get()).isNotEmpty().get().isEqualTo(queryId2);
+                assertThat(getPreparedCacheSize(session)).isEqualTo(1);
+              });
 
-      /* Okay, the latch triggered so cache processing should now be done.  Let's validate :allthethings: */
+      /* Okay, the type change and cache invalidation events should now be done. */
       assertThat(changedTypes.keySet()).isEqualTo(expectedChangedTypes);
       assertThat(removedQueryIds.get()).isNotEmpty().get().isEqualTo(queryId2);
       assertThat(getPreparedCacheSize(session)).isEqualTo(1);