Fixing potential race in ServerTimestamp tests

1 files changed, 8 insertions, 4 deletions

diff --git a/Firestore/Example/Tests/Integration/API/FIRServerTimestampTests.m b/Firestore/Example/Tests/Integration/API/FIRServerTimestampTests.m
index 5cda053..cc0ab29 100644
--- a/Firestore/Example/Tests/Integration/API/FIRServerTimestampTests.m
+++ b/Firestore/Example/Tests/Integration/API/FIRServerTimestampTests.m
@@ -101,15 +101,19 @@
 
 /** Waits for a snapshot with local writes. */
 - (FIRDocumentSnapshot *)waitForLocalEvent {
-  FIRDocumentSnapshot *snapshot = [_accumulator awaitEventWithName:@"Local event."];
-  XCTAssertTrue(snapshot.metadata.hasPendingWrites);
+  FIRDocumentSnapshot *snapshot;
+  do {
+    snapshot = [_accumulator awaitEventWithName:@"Local event."];
+  } while (!snapshot.metadata.hasPendingWrites);
   return snapshot;
 }
 
 /** Waits for a snapshot that has no pending writes */
 - (FIRDocumentSnapshot *)waitForRemoteEvent {
-  FIRDocumentSnapshot *snapshot = [_accumulator awaitEventWithName:@"Remote event."];
-  XCTAssertFalse(snapshot.metadata.hasPendingWrites);
+  FIRDocumentSnapshot *snapshot;
+  do {
+    snapshot = [_accumulator awaitEventWithName:@"Remote event."];
+  } while (snapshot.metadata.hasPendingWrites);
   return snapshot;
 }