test(datastore): Migrate datastore tests to use awaitility

java-datastore/google-cloud-datastore/pom.xml

@@ -245,6 +245,12 @@
       <version>0.15.0</version>
       <scope>test</scope>
     </dependency>
+    <dependency>
+      <groupId>org.awaitility</groupId>
+      <artifactId>awaitility</artifactId>
+      <version>4.3.0</version>
+      <scope>test</scope>
+    </dependency>
     <dependency>
       <groupId>com.google.cloud</groupId>
       <artifactId>google-cloud-trace</artifactId>

java-datastore/google-cloud-datastore/src/test/java/com/google/cloud/datastore/it/AbstractITDatastoreTest.java

@@ -26,6 +26,7 @@
 import static com.google.cloud.datastore.aggregation.Aggregation.sum;
 import static com.google.common.collect.Iterables.getOnlyElement;
 import static com.google.common.truth.Truth.assertThat;
+import static org.awaitility.Awaitility.await;
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertFalse;
 import static org.junit.Assert.assertNotEquals;
@@ -79,6 +80,7 @@
 import com.google.cloud.datastore.models.ExplainMetrics;
 import com.google.cloud.datastore.models.ExplainOptions;
 import com.google.cloud.datastore.models.PlanSummary;
+import com.google.cloud.testing.junit4.MultipleAttemptsRule;
 import com.google.common.base.Preconditions;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.Range;
@@ -97,8 +99,10 @@
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.Executors;
 import java.util.concurrent.Future;
+import java.util.concurrent.atomic.AtomicReference;
 import java.util.function.BiConsumer;
 import java.util.function.Consumer;
+import org.awaitility.core.ConditionTimeoutException;
 import org.junit.After;
 import org.junit.Before;
 import org.junit.Rule;
@@ -282,21 +286,28 @@ private <T> Iterator<T> getStronglyConsistentResults(Query scQuery, Query query)
     QueryResults<T> scResults = datastore.run(scQuery);
     List<T> scResultsCopy = makeResultsCopy(scResults);
     Set<T> scResultsSet = new HashSet<>(scResultsCopy);
-    int maxAttempts = 20;
-
-    while (maxAttempts > 0) {
-      --maxAttempts;
-      QueryResults<T> results = datastore.run(query);
-      List<T> resultsCopy = makeResultsCopy(results);
-      Set<T> resultsSet = new HashSet<>(resultsCopy);
-      if (scResultsSet.size() == resultsSet.size() && scResultsSet.containsAll(resultsSet)) {
-        return resultsCopy.iterator();
-      }
-      Thread.sleep(500);
+    AtomicReference<List<T>> finalResults = new AtomicReference<>();
+    try {
+      await()
+          .atMost(Duration.ofSeconds(10))
+          .pollInterval(Duration.ofMillis(500))
+          .until(
+              () -> {
+                QueryResults<T> results = datastore.run(query);
+                List<T> resultsCopy = makeResultsCopy(results);
+                Set<T> resultsSet = new HashSet<>(resultsCopy);
+                if (scResultsSet.size() == resultsSet.size()
+                    && scResultsSet.containsAll(resultsSet)) {
+                  finalResults.set(resultsCopy);
+                  return true;
+                }
+                return false;
+              });
+    } catch (ConditionTimeoutException e) {
+      throw new RuntimeException(
+          "reached max number of attempts to get strongly consistent results.", e);
     }
-
-    throw new RuntimeException(
-        "reached max number of attempts to get strongly consistent results.");
+    return finalResults.get().iterator();
   }
 
   private <T> List<T> makeResultsCopy(QueryResults<T> scResults) {
@@ -1331,8 +1342,12 @@ private void testCountAggregationReadTimeWith(Consumer<AggregationQuery.Builder>
             .build();
 
     datastore.put(entity1, entity2);
+    // Sleep to ensure time passes so that the subsequent Timestamp.now() is strictly
+    // after the commit time of entity1 and entity2.
     Thread.sleep(1000);
     Timestamp now = Timestamp.now();
+    // Sleep to ensure time passes so that the subsequent write of entity3 is strictly
+    // after 'now'. This allows testing ReadOption.readTime(now) deterministically.
     Thread.sleep(1000);
     datastore.put(entity3);
 
@@ -1784,8 +1799,12 @@ public void testGetWithReadTime() throws InterruptedException {
     try {
       datastore.put(Entity.newBuilder(key).set("str", "old_str_value").build());
 
+      // Sleep to ensure time passes so that the subsequent Timestamp.now() is strictly
+      // after the commit time of the old entity value.
       Thread.sleep(1000);
       Timestamp now = Timestamp.now();
+      // Sleep to ensure time passes so that the subsequent write of the new entity value
+      // is strictly after 'now'. This allows testing ReadOption.readTime(now) deterministically.
       Thread.sleep(1000);
 
       datastore.put(Entity.newBuilder(key).set("str", "new_str_value").build());
@@ -2102,8 +2121,12 @@ public void testQueryWithReadTime() throws InterruptedException {
             .build();
 
     datastore.put(entity1, entity2);
+    // Sleep to ensure time passes so that the subsequent Timestamp.now() is strictly
+    // after the commit time of entity1 and entity2.
     Thread.sleep(1000);
     Timestamp now = Timestamp.now();
+    // Sleep to ensure time passes so that the subsequent write of entity3 is strictly
+    // after 'now'. This allows testing ReadOption.readTime(now) deterministically.
     Thread.sleep(1000);
     datastore.put(entity3);
 

java-datastore/google-cloud-datastore/src/test/java/com/google/cloud/datastore/it/ITE2ETracingTest.java

@@ -31,6 +31,7 @@
 import static com.google.cloud.datastore.telemetry.TraceUtil.SPAN_NAME_TRANSACTION_RUN_QUERY;
 import static com.google.common.truth.Truth.assertThat;
 import static io.opentelemetry.semconv.resource.attributes.ResourceAttributes.SERVICE_NAME;
+import static org.awaitility.Awaitility.await;
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertFalse;
 import static org.junit.Assert.assertNotNull;
@@ -66,6 +67,7 @@
 import com.google.devtools.cloudtrace.v1.TraceSpan;
 import com.google.testing.junit.testparameterinjector.TestParameter;
 import com.google.testing.junit.testparameterinjector.TestParameterInjector;
+import io.grpc.StatusRuntimeException;
 import io.opentelemetry.api.GlobalOpenTelemetry;
 import io.opentelemetry.api.trace.Span;
 import io.opentelemetry.api.trace.SpanContext;
@@ -81,6 +83,7 @@
 import io.opentelemetry.sdk.trace.export.BatchSpanProcessor;
 import io.opentelemetry.sdk.trace.samplers.Sampler;
 import java.io.IOException;
+import java.time.Duration;
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.Collections;
@@ -89,8 +92,10 @@
 import java.util.Random;
 import java.util.TreeMap;
 import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicReference;
 import java.util.logging.Level;
 import java.util.logging.Logger;
+import org.awaitility.core.ConditionTimeoutException;
 import org.junit.After;
 import org.junit.AfterClass;
 import org.junit.Before;
@@ -452,44 +457,38 @@ protected void waitForTracesToComplete() throws Exception {
   protected void fetchAndValidateTrace(
       String traceId, int numExpectedSpans, List<List<String>> callStackList)
       throws InterruptedException {
-    // Large enough count to accommodate eventually consistent Cloud Trace backend
-    int numRetries = GET_TRACE_RETRY_COUNT;
     // Account for rootSpanName
-    numExpectedSpans++;
-
-    // Fetch traces
-    do {
-      try {
-        retrievedTrace = traceClient.getTrace(projectId, traceId);
-        assertEquals(traceId, retrievedTrace.getTraceId());
-
-        logger.info(
-            "expectedSpanCount="
-                + numExpectedSpans
-                + ", retrievedSpanCount="
-                + retrievedTrace.getSpansCount());
-      } catch (NotFoundException | DeadlineExceededException e) {
-        logger.info(
-            "Trace not found or deadline exceeded, retrying in "
-                + GET_TRACE_RETRY_BACKOFF_MILLIS
-                + " ms");
-      } catch (IndexOutOfBoundsException outOfBoundsException) {
-        logger.info("Call stack not found in trace. Retrying.");
-      }
-      if (retrievedTrace == null || numExpectedSpans != retrievedTrace.getSpansCount()) {
-        Thread.sleep(GET_TRACE_RETRY_BACKOFF_MILLIS);
-      }
-    } while (numRetries-- > 0
-        && (retrievedTrace == null || numExpectedSpans != retrievedTrace.getSpansCount()));
-
-    if (retrievedTrace == null || numExpectedSpans != retrievedTrace.getSpansCount()) {
+    final int expectedSpanCount = numExpectedSpans + 1;
+
+    try {
+      await()
+          .atMost(Duration.ofMillis((long) GET_TRACE_RETRY_COUNT * GET_TRACE_RETRY_BACKOFF_MILLIS))
+          .pollInterval(Duration.ofMillis(GET_TRACE_RETRY_BACKOFF_MILLIS))
+          .ignoreExceptionsInstanceOf(NotFoundException.class)
+          .ignoreExceptionsInstanceOf(DeadlineExceededException.class)
+          .ignoreExceptionsInstanceOf(IndexOutOfBoundsException.class)
+          .ignoreExceptionsInstanceOf(StatusRuntimeException.class)
+          .until(
+              () -> {
+                retrievedTrace = traceClient.getTrace(projectId, traceId);
+                assertEquals(traceId, retrievedTrace.getTraceId());
+                logger.info(
+                    "expectedSpanCount="
+                        + expectedSpanCount
+                        + ", retrievedSpanCount="
+                        + retrievedTrace.getSpansCount());
+                return retrievedTrace != null
+                    && expectedSpanCount == retrievedTrace.getSpansCount();
+              });
+    } catch (ConditionTimeoutException e) {
       throw new RuntimeException(
           "Expected number of spans: "
-              + numExpectedSpans
+              + expectedSpanCount
               + ", Actual number of spans: "
               + (retrievedTrace != null
                   ? retrievedTrace.getSpansList().toString()
-                  : "Trace NOT_FOUND"));
+                  : "Trace NOT_FOUND"),
+          e);
     }
 
     TraceContainer traceContainer = new TraceContainer(rootSpanName, retrievedTrace);
@@ -541,27 +540,33 @@ public void traceContainerTest() throws Exception {
     }
     waitForTracesToComplete();
 
-    Trace traceResp = null;
+    AtomicReference<Trace> traceRespHolder = new AtomicReference<>();
     int expectedSpanCount = 2;
 
-    int numRetries = GET_TRACE_RETRY_COUNT;
-    do {
-      try {
-        traceResp = traceClient.getTrace(projectId, customSpanContext.getTraceId());
-        if (traceResp.getSpansCount() == expectedSpanCount) {
-          logger.info("Success: Got " + expectedSpanCount + " spans.");
-          break;
-        }
-      } catch (NotFoundException notFoundException) {
-        Thread.sleep(GET_TRACE_RETRY_BACKOFF_MILLIS);
-        logger.info("Trace not found, retrying in " + GET_TRACE_RETRY_BACKOFF_MILLIS + " ms");
-      }
-      logger.info(
-          "Trace Found. The trace did not contain "
-              + expectedSpanCount
-              + " spans. Going to retry.");
-      numRetries--;
-    } while (numRetries > 0);
+    try {
+      await()
+          .atMost(Duration.ofMillis((long) GET_TRACE_RETRY_COUNT * GET_TRACE_RETRY_BACKOFF_MILLIS))
+          .pollInterval(Duration.ofMillis(GET_TRACE_RETRY_BACKOFF_MILLIS))
+          .ignoreExceptionsInstanceOf(NotFoundException.class)
+          .ignoreExceptionsInstanceOf(StatusRuntimeException.class)
+          .until(
+              () -> {
+                Trace trace = traceClient.getTrace(projectId, customSpanContext.getTraceId());
+                traceRespHolder.set(trace);
+                if (trace.getSpansCount() == expectedSpanCount) {
+                  logger.info("Success: Got " + expectedSpanCount + " spans.");
+                  return true;
+                }
+                logger.info(
+                    "Trace Found. The trace did not contain "
+                        + expectedSpanCount
+                        + " spans. Going to retry.");
+                return false;
+              });
+    } catch (ConditionTimeoutException ignored) {
+      // Ignore to let assertions below run and fail with descriptive messages
+    }
+    Trace traceResp = traceRespHolder.get();
 
     // Make sure we got as many spans as we expected.
     assertNotNull(traceResp);