Backport array contains / transform improvements from Web. (#1242)

5 files changed, 84 insertions, 85 deletions

diff --git a/Firestore/Example/Tests/Integration/API/FIRArrayTransformTests.mm b/Firestore/Example/Tests/Integration/API/FIRArrayTransformTests.mm
index 1c82461..0a5b413 100644
--- a/Firestore/Example/Tests/Integration/API/FIRArrayTransformTests.mm
+++ b/Firestore/Example/Tests/Integration/API/FIRArrayTransformTests.mm
@@ -36,7 +36,7 @@
   FIRDocumentReference *_docRef;
 
   // Accumulator used to capture events during the test.
-  FSTEventAccumulator *_accumulator;
+  FSTEventAccumulator<FIRDocumentSnapshot *> *_accumulator;
 
   // Listener registration for a listener maintained during the course of the test.
   id<FIRListenerRegistration> _listenerRegistration;
@@ -64,29 +64,11 @@
 
 #pragma mark - Test Helpers
 
-/** Waits for a snapshot with local writes. */
-- (FIRDocumentSnapshot *)waitForLocalEvent {
-  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;
-  do {
-    snapshot = [_accumulator awaitEventWithName:@"Remote event."];
-  } while (snapshot.metadata.hasPendingWrites);
-  return snapshot;
-}
-
 /** Writes some initial data and consumes the events generated. */
 - (void)writeInitialData:(NSDictionary<NSString *, id> *)data {
   [self writeDocumentRef:_docRef data:data];
-  XCTAssertEqualObjects([self waitForLocalEvent].data, data);
-  XCTAssertEqualObjects([self waitForRemoteEvent].data, data);
+  XCTAssertEqualObjects([_accumulator awaitLocalEvent].data, data);
+  XCTAssertEqualObjects([_accumulator awaitRemoteEvent].data, data);
 }
 
 #pragma mark - Test Cases
@@ -97,8 +79,8 @@
                       @"array" : [FIRFieldValue fieldValueForArrayUnion:@[ @1, @2 ]]
                     }];
   id expected = @{ @"array" : @[ @1, @2 ] };
-  XCTAssertEqualObjects([self waitForLocalEvent].data, expected);
-  XCTAssertEqualObjects([self waitForRemoteEvent].data, expected);
+  XCTAssertEqualObjects([_accumulator awaitLocalEvent].data, expected);
+  XCTAssertEqualObjects([_accumulator awaitRemoteEvent].data, expected);
 }
 
 - (void)testAppendToArrayViaUpdate {
@@ -110,8 +92,8 @@
                      }];
 
   id expected = @{ @"array" : @[ @1, @3, @2, @4 ] };
-  XCTAssertEqualObjects([self waitForLocalEvent].data, expected);
-  XCTAssertEqualObjects([self waitForRemoteEvent].data, expected);
+  XCTAssertEqualObjects([_accumulator awaitLocalEvent].data, expected);
+  XCTAssertEqualObjects([_accumulator awaitRemoteEvent].data, expected);
 }
 
 - (void)testAppendToArrayViaMergeSet {
@@ -123,8 +105,8 @@
                     }];
 
   id expected = @{ @"array" : @[ @1, @3, @2, @4 ] };
-  XCTAssertEqualObjects([self waitForLocalEvent].data, expected);
-  XCTAssertEqualObjects([self waitForRemoteEvent].data, expected);
+  XCTAssertEqualObjects([_accumulator awaitLocalEvent].data, expected);
+  XCTAssertEqualObjects([_accumulator awaitRemoteEvent].data, expected);
 }
 
 - (void)testAppendObjectToArrayViaUpdate {
@@ -137,8 +119,8 @@
                      }];
 
   id expected = @{ @"array" : @[ @{@"a" : @"hi"}, @{@"a" : @"bye"} ] };
-  XCTAssertEqualObjects([self waitForLocalEvent].data, expected);
-  XCTAssertEqualObjects([self waitForRemoteEvent].data, expected);
+  XCTAssertEqualObjects([_accumulator awaitLocalEvent].data, expected);
+  XCTAssertEqualObjects([_accumulator awaitRemoteEvent].data, expected);
 }
 
 - (void)testRemoveFromArrayViaUpdate {
@@ -150,8 +132,8 @@
                      }];
 
   id expected = @{ @"array" : @[ @3, @3 ] };
-  XCTAssertEqualObjects([self waitForLocalEvent].data, expected);
-  XCTAssertEqualObjects([self waitForRemoteEvent].data, expected);
+  XCTAssertEqualObjects([_accumulator awaitLocalEvent].data, expected);
+  XCTAssertEqualObjects([_accumulator awaitRemoteEvent].data, expected);
 }
 
 - (void)testRemoveFromArrayViaMergeSet {
@@ -163,8 +145,8 @@
                     }];
 
   id expected = @{ @"array" : @[ @3, @3 ] };
-  XCTAssertEqualObjects([self waitForLocalEvent].data, expected);
-  XCTAssertEqualObjects([self waitForRemoteEvent].data, expected);
+  XCTAssertEqualObjects([_accumulator awaitLocalEvent].data, expected);
+  XCTAssertEqualObjects([_accumulator awaitRemoteEvent].data, expected);
 }
 
 - (void)testRemoveObjectFromArrayViaUpdate {
@@ -176,8 +158,8 @@
                      }];
 
   id expected = @{ @"array" : @[ @{@"a" : @"bye"} ] };
-  XCTAssertEqualObjects([self waitForLocalEvent].data, expected);
-  XCTAssertEqualObjects([self waitForRemoteEvent].data, expected);
+  XCTAssertEqualObjects([_accumulator awaitLocalEvent].data, expected);
+  XCTAssertEqualObjects([_accumulator awaitRemoteEvent].data, expected);
 }
 
 @end
diff --git a/Firestore/Example/Tests/Integration/API/FIRServerTimestampTests.mm b/Firestore/Example/Tests/Integration/API/FIRServerTimestampTests.mm
index e1ad3d2..6a755c6 100644
--- a/Firestore/Example/Tests/Integration/API/FIRServerTimestampTests.mm
+++ b/Firestore/Example/Tests/Integration/API/FIRServerTimestampTests.mm
@@ -90,24 +90,6 @@
   XCTAssertEqualObjects(initialDataSnap.data, _initialData);
 }
 
-/** Waits for a snapshot with local writes. */
-- (FIRDocumentSnapshot *)waitForLocalEvent {
-  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;
-  do {
-    snapshot = [_accumulator awaitEventWithName:@"Remote event."];
-  } while (snapshot.metadata.hasPendingWrites);
-  return snapshot;
-}
-
 /** Verifies a snapshot containing _setData but with NSNull for the timestamps. */
 - (void)verifyTimestampsAreNullInSnapshot:(FIRDocumentSnapshot *)snapshot {
   XCTAssertEqualObjects(snapshot.data, [self expectedDataWithTimestamp:[NSNull null]]);
@@ -170,41 +152,42 @@
 
 - (void)testServerTimestampsWorkViaSet {
   [self writeDocumentRef:_docRef data:_setData];
-  [self verifyTimestampsAreNullInSnapshot:[self waitForLocalEvent]];
-  [self verifySnapshotWithResolvedTimestamps:[self waitForRemoteEvent]];
+  [self verifyTimestampsAreNullInSnapshot:[_accumulator awaitLocalEvent]];
+  [self verifySnapshotWithResolvedTimestamps:[_accumulator awaitRemoteEvent]];
 }
 
 - (void)testServerTimestampsWorkViaUpdate {
   [self writeInitialData];
   [self updateDocumentRef:_docRef data:_updateData];
-  [self verifyTimestampsAreNullInSnapshot:[self waitForLocalEvent]];
-  [self verifySnapshotWithResolvedTimestamps:[self waitForRemoteEvent]];
+  [self verifyTimestampsAreNullInSnapshot:[_accumulator awaitLocalEvent]];
+  [self verifySnapshotWithResolvedTimestamps:[_accumulator awaitRemoteEvent]];
 }
 
 - (void)testServerTimestampsWithEstimatedValue {
   [self writeDocumentRef:_docRef data:_setData];
-  [self verifyTimestampsAreEstimatedInSnapshot:[self waitForLocalEvent]];
-  [self verifySnapshotWithResolvedTimestamps:[self waitForRemoteEvent]];
+  [self verifyTimestampsAreEstimatedInSnapshot:[_accumulator awaitLocalEvent]];
+  [self verifySnapshotWithResolvedTimestamps:[_accumulator awaitRemoteEvent]];
 }
 
 - (void)testServerTimestampsWithPreviousValue {
   [self writeDocumentRef:_docRef data:_setData];
-  [self verifyTimestampsInSnapshot:[self waitForLocalEvent] fromPreviousSnapshot:nil];
-  FIRDocumentSnapshot *remoteSnapshot = [self waitForRemoteEvent];
+  [self verifyTimestampsInSnapshot:[_accumulator awaitLocalEvent] fromPreviousSnapshot:nil];
+  FIRDocumentSnapshot *remoteSnapshot = [_accumulator awaitRemoteEvent];
 
   [_docRef updateData:_updateData];
-  [self verifyTimestampsInSnapshot:[self waitForLocalEvent] fromPreviousSnapshot:remoteSnapshot];
+  [self verifyTimestampsInSnapshot:[_accumulator awaitLocalEvent]
+              fromPreviousSnapshot:remoteSnapshot];
 
-  [self verifySnapshotWithResolvedTimestamps:[self waitForRemoteEvent]];
+  [self verifySnapshotWithResolvedTimestamps:[_accumulator awaitRemoteEvent]];
 }
 
 - (void)testServerTimestampsWithPreviousValueOfDifferentType {
   [self writeDocumentRef:_docRef data:_setData];
-  [self verifyTimestampsInSnapshot:[self waitForLocalEvent] fromPreviousSnapshot:nil];
-  [self verifySnapshotWithResolvedTimestamps:[self waitForRemoteEvent]];
+  [self verifyTimestampsInSnapshot:[_accumulator awaitLocalEvent] fromPreviousSnapshot:nil];
+  [self verifySnapshotWithResolvedTimestamps:[_accumulator awaitRemoteEvent]];
 
   [_docRef updateData:@{@"a" : [FIRFieldValue fieldValueForServerTimestamp]}];
-  FIRDocumentSnapshot *localSnapshot = [self waitForLocalEvent];
+  FIRDocumentSnapshot *localSnapshot = [_accumulator awaitLocalEvent];
   XCTAssertEqualObjects([localSnapshot valueForField:@"a"], [NSNull null]);
   XCTAssertEqualObjects(
       [localSnapshot valueForField:@"a" serverTimestampBehavior:FIRServerTimestampBehaviorPrevious],
@@ -213,7 +196,7 @@
       [[localSnapshot valueForField:@"a" serverTimestampBehavior:FIRServerTimestampBehaviorEstimate]
           isKindOfClass:[FIRTimestamp class]]);
 
-  FIRDocumentSnapshot *remoteSnapshot = [self waitForRemoteEvent];
+  FIRDocumentSnapshot *remoteSnapshot = [_accumulator awaitRemoteEvent];
   XCTAssertTrue([[remoteSnapshot valueForField:@"a"] isKindOfClass:[FIRTimestamp class]]);
   XCTAssertTrue([
       [remoteSnapshot valueForField:@"a" serverTimestampBehavior:FIRServerTimestampBehaviorPrevious]
@@ -225,55 +208,55 @@
 
 - (void)testServerTimestampsWithConsecutiveUpdates {
   [self writeDocumentRef:_docRef data:_setData];
-  [self verifyTimestampsInSnapshot:[self waitForLocalEvent] fromPreviousSnapshot:nil];
-  [self verifySnapshotWithResolvedTimestamps:[self waitForRemoteEvent]];
+  [self verifyTimestampsInSnapshot:[_accumulator awaitLocalEvent] fromPreviousSnapshot:nil];
+  [self verifySnapshotWithResolvedTimestamps:[_accumulator awaitRemoteEvent]];
 
   [self disableNetwork];
 
   [_docRef updateData:@{@"a" : [FIRFieldValue fieldValueForServerTimestamp]}];
-  FIRDocumentSnapshot *localSnapshot = [self waitForLocalEvent];
+  FIRDocumentSnapshot *localSnapshot = [_accumulator awaitLocalEvent];
   XCTAssertEqualObjects(
       [localSnapshot valueForField:@"a" serverTimestampBehavior:FIRServerTimestampBehaviorPrevious],
       @42);
 
   [_docRef updateData:@{@"a" : [FIRFieldValue fieldValueForServerTimestamp]}];
-  localSnapshot = [self waitForLocalEvent];
+  localSnapshot = [_accumulator awaitLocalEvent];
   XCTAssertEqualObjects(
       [localSnapshot valueForField:@"a" serverTimestampBehavior:FIRServerTimestampBehaviorPrevious],
       @42);
 
   [self enableNetwork];
 
-  FIRDocumentSnapshot *remoteSnapshot = [self waitForRemoteEvent];
+  FIRDocumentSnapshot *remoteSnapshot = [_accumulator awaitRemoteEvent];
   XCTAssertTrue([[remoteSnapshot valueForField:@"a"] isKindOfClass:[FIRTimestamp class]]);
 }
 
 - (void)testServerTimestampsPreviousValueFromLocalMutation {
   [self writeDocumentRef:_docRef data:_setData];
-  [self verifyTimestampsInSnapshot:[self waitForLocalEvent] fromPreviousSnapshot:nil];
-  [self verifySnapshotWithResolvedTimestamps:[self waitForRemoteEvent]];
+  [self verifyTimestampsInSnapshot:[_accumulator awaitLocalEvent] fromPreviousSnapshot:nil];
+  [self verifySnapshotWithResolvedTimestamps:[_accumulator awaitRemoteEvent]];
 
   [self disableNetwork];
 
   [_docRef updateData:@{@"a" : [FIRFieldValue fieldValueForServerTimestamp]}];
-  FIRDocumentSnapshot *localSnapshot = [self waitForLocalEvent];
+  FIRDocumentSnapshot *localSnapshot = [_accumulator awaitLocalEvent];
   XCTAssertEqualObjects(
       [localSnapshot valueForField:@"a" serverTimestampBehavior:FIRServerTimestampBehaviorPrevious],
       @42);
 
   [_docRef updateData:@{ @"a" : @1337 }];
-  localSnapshot = [self waitForLocalEvent];
+  localSnapshot = [_accumulator awaitLocalEvent];
   XCTAssertEqualObjects([localSnapshot valueForField:@"a"], @1337);
 
   [_docRef updateData:@{@"a" : [FIRFieldValue fieldValueForServerTimestamp]}];
-  localSnapshot = [self waitForLocalEvent];
+  localSnapshot = [_accumulator awaitLocalEvent];
   XCTAssertEqualObjects(
       [localSnapshot valueForField:@"a" serverTimestampBehavior:FIRServerTimestampBehaviorPrevious],
       @1337);
 
   [self enableNetwork];
 
-  FIRDocumentSnapshot *remoteSnapshot = [self waitForRemoteEvent];
+  FIRDocumentSnapshot *remoteSnapshot = [_accumulator awaitRemoteEvent];
   XCTAssertTrue([[remoteSnapshot valueForField:@"a"] isKindOfClass:[FIRTimestamp class]]);
 }
 
@@ -282,7 +265,7 @@
     [transaction setData:_setData forDocument:_docRef];
   }];
 
-  [self verifySnapshotWithResolvedTimestamps:[self waitForRemoteEvent]];
+  [self verifySnapshotWithResolvedTimestamps:[_accumulator awaitRemoteEvent]];
 }
 
 - (void)testServerTimestampsWorkViaTransactionUpdate {
@@ -290,7 +273,7 @@
   [self runTransactionBlock:^(FIRTransaction *transaction) {
     [transaction updateData:_updateData forDocument:_docRef];
   }];
-  [self verifySnapshotWithResolvedTimestamps:[self waitForRemoteEvent]];
+  [self verifySnapshotWithResolvedTimestamps:[_accumulator awaitRemoteEvent]];
 }
 
 - (void)testServerTimestampsFailViaUpdateOnNonexistentDocument {
diff --git a/Firestore/Example/Tests/Integration/API/FIRValidationTests.mm b/Firestore/Example/Tests/Integration/API/FIRValidationTests.mm
index 6d10aba..8af8d15 100644
--- a/Firestore/Example/Tests/Integration/API/FIRValidationTests.mm
+++ b/Firestore/Example/Tests/Integration/API/FIRValidationTests.mm
@@ -507,8 +507,8 @@
   FSTAssertThrows([collection queryWhereFieldPath:[FIRFieldPath documentID] isEqualTo:@1], reason);
 
   reason =
-      @"Invalid query. You can't do arrayContains queries on document ID since document IDs are "
-      @"not arrays.";
+      @"Invalid query. You can't perform arrayContains queries on document ID since document IDs "
+       "are not arrays.";
   FSTAssertThrows([collection queryWhereFieldPath:[FIRFieldPath documentID] arrayContains:@1],
                   reason);
 }
diff --git a/Firestore/Example/Tests/Util/FSTEventAccumulator.h b/Firestore/Example/Tests/Util/FSTEventAccumulator.h
index baa501b..58b802b 100644
--- a/Firestore/Example/Tests/Util/FSTEventAccumulator.h
+++ b/Firestore/Example/Tests/Util/FSTEventAccumulator.h
@@ -25,15 +25,21 @@ NS_ASSUME_NONNULL_BEGIN
 
 typedef void (^FSTValueEventHandler)(id _Nullable, NSError *_Nullable error);
 
-@interface FSTEventAccumulator : NSObject
+@interface FSTEventAccumulator <EventType> : NSObject
 
 + (instancetype)accumulatorForTest:(XCTestCase *)testCase;
 
 - (instancetype)init NS_UNAVAILABLE;
 
-- (id)awaitEventWithName:(NSString *)name;
+- (EventType)awaitEventWithName:(NSString *)name;
 
-- (NSArray<id> *)awaitEvents:(NSUInteger)events name:(NSString *)name;
+- (NSArray<EventType> *)awaitEvents:(NSUInteger)events name:(NSString *)name;
+
+/** Waits for a latency compensated local snapshot. */
+- (EventType)awaitLocalEvent;
+
+/** Waits for a snapshot that has no pending writes */
+- (EventType)awaitRemoteEvent;
 
 @property(nonatomic, strong, readonly) FSTValueEventHandler valueEventHandler;
 
diff --git a/Firestore/Example/Tests/Util/FSTEventAccumulator.mm b/Firestore/Example/Tests/Util/FSTEventAccumulator.mm
index 623ba2d..3ab6035 100644
--- a/Firestore/Example/Tests/Util/FSTEventAccumulator.mm
+++ b/Firestore/Example/Tests/Util/FSTEventAccumulator.mm
@@ -18,6 +18,9 @@
 
 #import <XCTest/XCTest.h>
 
+#import "Firestore/Source/Public/FIRDocumentSnapshot.h"
+#import "Firestore/Source/Public/FIRQuerySnapshot.h"
+#import "Firestore/Source/Public/FIRSnapshotMetadata.h"
 #import "Firestore/Source/Util/FSTAssert.h"
 
 #import "Firestore/Example/Tests/Util/XCTestCase+Await.h"
@@ -68,6 +71,31 @@ NS_ASSUME_NONNULL_BEGIN
   return events[0];
 }
 
+- (id)awaitLocalEvent {
+  id event;
+  do {
+    event = [self awaitEventWithName:@"Local Event"];
+  } while (![self hasPendingWrites:event]);
+  return event;
+}
+
+- (id)awaitRemoteEvent {
+  id event;
+  do {
+    event = [self awaitEventWithName:@"Remote Event"];
+  } while ([self hasPendingWrites:event]);
+  return event;
+}
+
+- (BOOL)hasPendingWrites:(id)event {
+  if ([event isKindOfClass:[FIRDocumentSnapshot class]]) {
+    return ((FIRDocumentSnapshot *)event).metadata.hasPendingWrites;
+  } else {
+    FSTAssert([event isKindOfClass:[FIRQuerySnapshot class]], @"Unexpected event: %@", event);
+    return ((FIRQuerySnapshot *)event).metadata.hasPendingWrites;
+  }
+}
+
 - (void (^)(id _Nullable, NSError *_Nullable))valueEventHandler {
   return ^void(id _Nullable value, NSError *_Nullable error) {
     // We can't store nil in the _events array, but these are still interesting to tests so store