Skip to content

Commit 2f36a71

Browse files
lqiu96lqiu96
authored
test(datastore): Migrate datastore tests to use awaitility (#13668)
Also removes Datastore's MultipleAttemptsRule to use the version from java-core. From: #13645 (comment)
1 parent 557d430 commit 2f36a71

5 files changed

Lines changed: 94 additions & 175 deletions

File tree

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

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -245,6 +245,12 @@
245245
<version>0.15.0</version>
246246
<scope>test</scope>
247247
</dependency>
248+
<dependency>
249+
<groupId>org.awaitility</groupId>
250+
<artifactId>awaitility</artifactId>
251+
<version>4.3.0</version>
252+
<scope>test</scope>
253+
</dependency>
248254
<dependency>
249255
<groupId>com.google.cloud</groupId>
250256
<artifactId>google-cloud-trace</artifactId>

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

Lines changed: 37 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -26,6 +26,7 @@
2626
import static com.google.cloud.datastore.aggregation.Aggregation.sum;
2727
import static com.google.common.collect.Iterables.getOnlyElement;
2828
import static com.google.common.truth.Truth.assertThat;
29+
import static org.awaitility.Awaitility.await;
2930
import static org.junit.Assert.assertEquals;
3031
import static org.junit.Assert.assertFalse;
3132
import static org.junit.Assert.assertNotEquals;
@@ -79,6 +80,7 @@
7980
import com.google.cloud.datastore.models.ExplainMetrics;
8081
import com.google.cloud.datastore.models.ExplainOptions;
8182
import com.google.cloud.datastore.models.PlanSummary;
83+
import com.google.cloud.testing.junit4.MultipleAttemptsRule;
8284
import com.google.common.base.Preconditions;
8385
import com.google.common.collect.ImmutableList;
8486
import com.google.common.collect.Range;
@@ -97,8 +99,10 @@
9799
import java.util.concurrent.ExecutorService;
98100
import java.util.concurrent.Executors;
99101
import java.util.concurrent.Future;
102+
import java.util.concurrent.atomic.AtomicReference;
100103
import java.util.function.BiConsumer;
101104
import java.util.function.Consumer;
105+
import org.awaitility.core.ConditionTimeoutException;
102106
import org.junit.After;
103107
import org.junit.Before;
104108
import org.junit.Rule;
@@ -282,21 +286,28 @@ private <T> Iterator<T> getStronglyConsistentResults(Query scQuery, Query query)
282286
QueryResults<T> scResults = datastore.run(scQuery);
283287
List<T> scResultsCopy = makeResultsCopy(scResults);
284288
Set<T> scResultsSet = new HashSet<>(scResultsCopy);
285-
int maxAttempts = 20;
286-
287-
while (maxAttempts > 0) {
288-
--maxAttempts;
289-
QueryResults<T> results = datastore.run(query);
290-
List<T> resultsCopy = makeResultsCopy(results);
291-
Set<T> resultsSet = new HashSet<>(resultsCopy);
292-
if (scResultsSet.size() == resultsSet.size() && scResultsSet.containsAll(resultsSet)) {
293-
return resultsCopy.iterator();
294-
}
295-
Thread.sleep(500);
289+
AtomicReference<List<T>> finalResults = new AtomicReference<>();
290+
try {
291+
await()
292+
.atMost(Duration.ofSeconds(10))
293+
.pollInterval(Duration.ofMillis(500))
294+
.until(
295+
() -> {
296+
QueryResults<T> results = datastore.run(query);
297+
List<T> resultsCopy = makeResultsCopy(results);
298+
Set<T> resultsSet = new HashSet<>(resultsCopy);
299+
if (scResultsSet.size() == resultsSet.size()
300+
&& scResultsSet.containsAll(resultsSet)) {
301+
finalResults.set(resultsCopy);
302+
return true;
303+
}
304+
return false;
305+
});
306+
} catch (ConditionTimeoutException e) {
307+
throw new RuntimeException(
308+
"reached max number of attempts to get strongly consistent results.", e);
296309
}
297-
298-
throw new RuntimeException(
299-
"reached max number of attempts to get strongly consistent results.");
310+
return finalResults.get().iterator();
300311
}
301312

302313
private <T> List<T> makeResultsCopy(QueryResults<T> scResults) {
@@ -1331,8 +1342,12 @@ private void testCountAggregationReadTimeWith(Consumer<AggregationQuery.Builder>
13311342
.build();
13321343

13331344
datastore.put(entity1, entity2);
1345+
// Sleep to ensure time passes so that the subsequent Timestamp.now() is strictly
1346+
// after the commit time of entity1 and entity2.
13341347
Thread.sleep(1000);
13351348
Timestamp now = Timestamp.now();
1349+
// Sleep to ensure time passes so that the subsequent write of entity3 is strictly
1350+
// after 'now'. This allows testing ReadOption.readTime(now) deterministically.
13361351
Thread.sleep(1000);
13371352
datastore.put(entity3);
13381353

@@ -1784,8 +1799,12 @@ public void testGetWithReadTime() throws InterruptedException {
17841799
try {
17851800
datastore.put(Entity.newBuilder(key).set("str", "old_str_value").build());
17861801

1802+
// Sleep to ensure time passes so that the subsequent Timestamp.now() is strictly
1803+
// after the commit time of the old entity value.
17871804
Thread.sleep(1000);
17881805
Timestamp now = Timestamp.now();
1806+
// Sleep to ensure time passes so that the subsequent write of the new entity value
1807+
// is strictly after 'now'. This allows testing ReadOption.readTime(now) deterministically.
17891808
Thread.sleep(1000);
17901809

17911810
datastore.put(Entity.newBuilder(key).set("str", "new_str_value").build());
@@ -2102,8 +2121,12 @@ public void testQueryWithReadTime() throws InterruptedException {
21022121
.build();
21032122

21042123
datastore.put(entity1, entity2);
2124+
// Sleep to ensure time passes so that the subsequent Timestamp.now() is strictly
2125+
// after the commit time of entity1 and entity2.
21052126
Thread.sleep(1000);
21062127
Timestamp now = Timestamp.now();
2128+
// Sleep to ensure time passes so that the subsequent write of entity3 is strictly
2129+
// after 'now'. This allows testing ReadOption.readTime(now) deterministically.
21072130
Thread.sleep(1000);
21082131
datastore.put(entity3);
21092132

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

Lines changed: 51 additions & 53 deletions
Original file line numberDiff line numberDiff line change
@@ -31,15 +31,14 @@
3131
import static com.google.cloud.datastore.telemetry.TraceUtil.SPAN_NAME_TRANSACTION_RUN_QUERY;
3232
import static com.google.common.truth.Truth.assertThat;
3333
import static io.opentelemetry.semconv.resource.attributes.ResourceAttributes.SERVICE_NAME;
34+
import static org.awaitility.Awaitility.await;
3435
import static org.junit.Assert.assertEquals;
3536
import static org.junit.Assert.assertFalse;
3637
import static org.junit.Assert.assertNotNull;
3738
import static org.junit.Assert.assertNull;
3839
import static org.junit.Assert.assertTrue;
3940

4041
import com.google.api.gax.core.FixedCredentialsProvider;
41-
import com.google.api.gax.rpc.DeadlineExceededException;
42-
import com.google.api.gax.rpc.NotFoundException;
4342
import com.google.auth.Credentials;
4443
import com.google.cloud.datastore.AggregationQuery;
4544
import com.google.cloud.datastore.AggregationResult;
@@ -81,6 +80,7 @@
8180
import io.opentelemetry.sdk.trace.export.BatchSpanProcessor;
8281
import io.opentelemetry.sdk.trace.samplers.Sampler;
8382
import java.io.IOException;
83+
import java.time.Duration;
8484
import java.util.ArrayList;
8585
import java.util.Arrays;
8686
import java.util.Collections;
@@ -89,8 +89,10 @@
8989
import java.util.Random;
9090
import java.util.TreeMap;
9191
import java.util.concurrent.TimeUnit;
92+
import java.util.concurrent.atomic.AtomicReference;
9293
import java.util.logging.Level;
9394
import java.util.logging.Logger;
95+
import org.awaitility.core.ConditionTimeoutException;
9496
import org.junit.After;
9597
import org.junit.AfterClass;
9698
import org.junit.Before;
@@ -474,44 +476,35 @@ protected void waitForTracesToComplete() throws Exception {
474476
protected void fetchAndValidateTrace(
475477
String traceId, int numExpectedSpans, List<List<String>> callStackList)
476478
throws InterruptedException {
477-
// Large enough count to accommodate eventually consistent Cloud Trace backend
478-
int numRetries = GET_TRACE_RETRY_COUNT;
479479
// Account for rootSpanName
480-
numExpectedSpans++;
481-
482-
// Fetch traces
483-
do {
484-
try {
485-
retrievedTrace = traceClient.getTrace(projectId, traceId);
486-
assertEquals(traceId, retrievedTrace.getTraceId());
487-
488-
logger.info(
489-
"expectedSpanCount="
490-
+ numExpectedSpans
491-
+ ", retrievedSpanCount="
492-
+ retrievedTrace.getSpansCount());
493-
} catch (NotFoundException | DeadlineExceededException e) {
494-
logger.info(
495-
"Trace not found or deadline exceeded, retrying in "
496-
+ GET_TRACE_RETRY_BACKOFF_MILLIS
497-
+ " ms");
498-
} catch (IndexOutOfBoundsException outOfBoundsException) {
499-
logger.info("Call stack not found in trace. Retrying.");
500-
}
501-
if (retrievedTrace == null || numExpectedSpans != retrievedTrace.getSpansCount()) {
502-
Thread.sleep(GET_TRACE_RETRY_BACKOFF_MILLIS);
503-
}
504-
} while (numRetries-- > 0
505-
&& (retrievedTrace == null || numExpectedSpans != retrievedTrace.getSpansCount()));
480+
final int expectedSpanCount = numExpectedSpans + 1;
506481

507-
if (retrievedTrace == null || numExpectedSpans != retrievedTrace.getSpansCount()) {
482+
try {
483+
await()
484+
.atMost(Duration.ofMillis((long) GET_TRACE_RETRY_COUNT * GET_TRACE_RETRY_BACKOFF_MILLIS))
485+
.pollInterval(Duration.ofMillis(GET_TRACE_RETRY_BACKOFF_MILLIS))
486+
.ignoreExceptionsInstanceOf(Throwable.class)
487+
.until(
488+
() -> {
489+
retrievedTrace = traceClient.getTrace(projectId, traceId);
490+
assertEquals(traceId, retrievedTrace.getTraceId());
491+
logger.info(
492+
"expectedSpanCount="
493+
+ expectedSpanCount
494+
+ ", retrievedSpanCount="
495+
+ retrievedTrace.getSpansCount());
496+
return retrievedTrace != null
497+
&& expectedSpanCount == retrievedTrace.getSpansCount();
498+
});
499+
} catch (ConditionTimeoutException e) {
508500
throw new RuntimeException(
509501
"Expected number of spans: "
510-
+ numExpectedSpans
502+
+ expectedSpanCount
511503
+ ", Actual number of spans: "
512504
+ (retrievedTrace != null
513505
? retrievedTrace.getSpansList().toString()
514-
: "Trace NOT_FOUND"));
506+
: "Trace NOT_FOUND"),
507+
e);
515508
}
516509

517510
TraceContainer traceContainer = new TraceContainer(rootSpanName, retrievedTrace);
@@ -527,6 +520,7 @@ protected void fetchAndValidateTrace(
527520
logger.info("Checking if TraceContainer contains the callStack");
528521
String[] expectedCallList = new String[expectedCallStack.size()];
529522
if (!traceContainer.containsCallStack(expectedCallStack.toArray(expectedCallList))) {
523+
logger.severe("CallStack not found in TraceContainer.");
530524
throw new RuntimeException(
531525
"Expected spans: "
532526
+ Arrays.toString(expectedCallList)
@@ -535,7 +529,6 @@ protected void fetchAndValidateTrace(
535529
? retrievedTrace.getSpansList().toString()
536530
: "Trace NOT_FOUND"));
537531
}
538-
logger.severe("CallStack not found in TraceContainer.");
539532
}
540533
}
541534

@@ -563,27 +556,32 @@ public void traceContainerTest() throws Exception {
563556
}
564557
waitForTracesToComplete();
565558

566-
Trace traceResp = null;
559+
AtomicReference<Trace> traceRespHolder = new AtomicReference<>();
567560
int expectedSpanCount = 2;
568561

569-
int numRetries = GET_TRACE_RETRY_COUNT;
570-
do {
571-
try {
572-
traceResp = traceClient.getTrace(projectId, customSpanContext.getTraceId());
573-
if (traceResp.getSpansCount() == expectedSpanCount) {
574-
logger.info("Success: Got " + expectedSpanCount + " spans.");
575-
break;
576-
}
577-
} catch (NotFoundException notFoundException) {
578-
Thread.sleep(GET_TRACE_RETRY_BACKOFF_MILLIS);
579-
logger.info("Trace not found, retrying in " + GET_TRACE_RETRY_BACKOFF_MILLIS + " ms");
580-
}
581-
logger.info(
582-
"Trace Found. The trace did not contain "
583-
+ expectedSpanCount
584-
+ " spans. Going to retry.");
585-
numRetries--;
586-
} while (numRetries > 0);
562+
try {
563+
await()
564+
.atMost(Duration.ofMillis((long) GET_TRACE_RETRY_COUNT * GET_TRACE_RETRY_BACKOFF_MILLIS))
565+
.pollInterval(Duration.ofMillis(GET_TRACE_RETRY_BACKOFF_MILLIS))
566+
.ignoreExceptionsInstanceOf(Throwable.class)
567+
.until(
568+
() -> {
569+
Trace trace = traceClient.getTrace(projectId, customSpanContext.getTraceId());
570+
traceRespHolder.set(trace);
571+
if (trace.getSpansCount() == expectedSpanCount) {
572+
logger.info("Success: Got " + expectedSpanCount + " spans.");
573+
return true;
574+
}
575+
logger.info(
576+
"Trace Found. The trace did not contain "
577+
+ expectedSpanCount
578+
+ " spans. Going to retry.");
579+
return false;
580+
});
581+
} catch (ConditionTimeoutException ignored) {
582+
// Ignore to let assertions below run and fail with descriptive messages
583+
}
584+
Trace traceResp = traceRespHolder.get();
587585

588586
// Make sure we got as many spans as we expected.
589587
assertNotNull(traceResp);

‎java-datastore/google-cloud-datastore/src/test/java/com/google/cloud/datastore/it/MultipleAttemptsRule.java‎

Lines changed: 0 additions & 71 deletions
This file was deleted.

0 commit comments

Comments
 (0)