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

2 files changed, 37 insertions, 3 deletions

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