Firestore: FIRQueryTests.mm: add a test that resumes a query with existence filter (#10974)

Author: dconeybe

Firestore/Example/Tests/Integration/API/FIRQueryTests.mm

@@ -22,6 +22,16 @@
 #import "Firestore/Example/Tests/Util/FSTHelpers.h"
 #import "Firestore/Example/Tests/Util/FSTIntegrationTestCase.h"
 
+namespace {
+
+NSArray<NSString *> *SortedStringsNotIn(NSSet<NSString *> *set, NSSet<NSString *> *remove) {
+  NSMutableSet<NSString *> *mutableSet = [NSMutableSet setWithSet:set];
+  [mutableSet minusSet:remove];
+  return [mutableSet.allObjects sortedArrayUsingSelector:@selector(caseInsensitiveCompare:)];
+}
+
+}  // namespace
+
 @interface FIRQueryTests : FSTIntegrationTestCase
 @end
 
@@ -1180,4 +1190,88 @@ - (void)testOrderByEquality {
                      matchesResult:@[ @"doc6", @"doc3" ]];
 }
 
+- (void)testResumingAQueryShouldUseExistenceFilterToDetectDeletes {
+  // Prepare the names and contents of the 100 documents to create.
+  NSMutableDictionary<NSString *, NSDictionary<NSString *, id> *> *testDocs =
+      [[NSMutableDictionary alloc] init];
+  for (int i = 0; i < 100; i++) {
+    [testDocs setValue:@{@"key" : @42} forKey:[NSString stringWithFormat:@"doc%@", @(1000 + i)]];
+  }
+
+  // Create 100 documents in a new collection.
+  FIRCollectionReference *collRef = [self collectionRefWithDocuments:testDocs];
+
+  // Run a query to populate the local cache with the 100 documents and a resume token.
+  FIRQuerySnapshot *querySnapshot1 = [self readDocumentSetForRef:collRef
+                                                          source:FIRFirestoreSourceDefault];
+  XCTAssertEqual(querySnapshot1.count, 100, @"querySnapshot1.count has an unexpected value");
+  NSArray<FIRDocumentReference *> *createdDocuments =
+      FIRDocumentReferenceArrayFromQuerySnapshot(querySnapshot1);
+
+  // Delete 50 of the 100 documents. Do this in a transaction, rather than
+  // [FIRDocumentReference deleteDocument], to avoid affecting the local cache.
+  NSSet<NSString *> *deletedDocumentIds;
+  {
+    NSMutableArray<NSString *> *deletedDocumentIdsAccumulator = [[NSMutableArray alloc] init];
+    XCTestExpectation *expectation = [self expectationWithDescription:@"DeleteTransaction"];
+    [collRef.firestore
+        runTransactionWithBlock:^id _Nullable(FIRTransaction *transaction, NSError **) {
+          for (decltype(createdDocuments.count) i = 0; i < createdDocuments.count; i += 2) {
+            FIRDocumentReference *documentToDelete = createdDocuments[i];
+            [transaction deleteDocument:documentToDelete];
+            [deletedDocumentIdsAccumulator addObject:documentToDelete.documentID];
+          }
+          return @"document deletion successful";
+        }
+        completion:^(id, NSError *) {
+          [expectation fulfill];
+        }];
+    [self awaitExpectation:expectation];
+    deletedDocumentIds = [NSSet setWithArray:deletedDocumentIdsAccumulator];
+  }
+  XCTAssertEqual(deletedDocumentIds.count, 50u, @"deletedDocumentIds has the wrong size");
+
+  // Wait for 10 seconds, during which Watch will stop tracking the query and will send an existence
+  // filter rather than "delete" events when the query is resumed.
+  [NSThread sleepForTimeInterval:10.0f];
+
+  // Resume the query and save the resulting snapshot for verification.
+  FIRQuerySnapshot *querySnapshot2 = [self readDocumentSetForRef:collRef
+                                                          source:FIRFirestoreSourceDefault];
+
+  // Verify that the snapshot from the resumed query contains the expected documents; that is,
+  // that it contains the 50 documents that were _not_ deleted.
+  // TODO(b/270731363): Remove the "if" condition below once the Firestore Emulator is fixed to
+  // send an existence filter. At the time of writing, the Firestore emulator fails to send an
+  // existence filter, resulting in the client including the deleted documents in the snapshot
+  // of the resumed query.
+  if (!([FSTIntegrationTestCase isRunningAgainstEmulator] && querySnapshot2.count == 100)) {
+    NSSet<NSString *> *actualDocumentIds =
+        [NSSet setWithArray:FIRQuerySnapshotGetIDs(querySnapshot2)];
+    NSSet<NSString *> *expectedDocumentIds;
+    {
+      NSMutableArray<NSString *> *expectedDocumentIdsAccumulator = [[NSMutableArray alloc] init];
+      for (FIRDocumentReference *documentRef in createdDocuments) {
+        if (![deletedDocumentIds containsObject:documentRef.documentID]) {
+          [expectedDocumentIdsAccumulator addObject:documentRef.documentID];
+        }
+      }
+      expectedDocumentIds = [NSSet setWithArray:expectedDocumentIdsAccumulator];
+    }
+    if (![actualDocumentIds isEqualToSet:expectedDocumentIds]) {
+      NSArray<NSString *> *unexpectedDocumentIds =
+          SortedStringsNotIn(actualDocumentIds, expectedDocumentIds);
+      NSArray<NSString *> *missingDocumentIds =
+          SortedStringsNotIn(expectedDocumentIds, actualDocumentIds);
+      XCTFail(@"querySnapshot2 contained %lu documents (expected %lu): "
+              @"%lu unexpected and %lu missing; "
+              @"unexpected documents: %@; missing documents: %@",
+              (unsigned long)actualDocumentIds.count, (unsigned long)expectedDocumentIds.count,
+              (unsigned long)unexpectedDocumentIds.count, (unsigned long)missingDocumentIds.count,
+              [unexpectedDocumentIds componentsJoinedByString:@", "],
+              [missingDocumentIds componentsJoinedByString:@", "]);
+    }
+  }
+}
+
 @end

Firestore/Example/Tests/Util/FSTIntegrationTestCase.h

@@ -142,6 +142,9 @@ NSArray<NSString *> *FIRQuerySnapshotGetIDs(FIRQuerySnapshot *docs);
  * in order of { type, doc title, doc data }. */
 NSArray<NSArray<id> *> *FIRQuerySnapshotGetDocChangesData(FIRQuerySnapshot *docs);
 
+/** Gets the FIRDocumentReference objects from a FIRQuerySnapshot and returns them. */
+NSArray<FIRDocumentReference *> *FIRDocumentReferenceArrayFromQuerySnapshot(FIRQuerySnapshot *);
+
 #if __cplusplus
 }  // extern "C"
 #endif

Firestore/Example/Tests/Util/FSTIntegrationTestCase.mm

@@ -25,6 +25,7 @@
 #import <FirebaseFirestore/FIRQuerySnapshot.h>
 #import <FirebaseFirestore/FIRSnapshotMetadata.h>
 #import <FirebaseFirestore/FIRTransaction.h>
+#import <FirebaseFirestore/FIRWriteBatch.h>
 
 #include <exception>
 #include <memory>
@@ -393,11 +394,48 @@ - (FIRCollectionReference *)collectionRefWithDocuments:
 
 - (void)writeAllDocuments:(NSDictionary<NSString *, NSDictionary<NSString *, id> *> *)documents
              toCollection:(FIRCollectionReference *)collection {
+  NSMutableArray<XCTestExpectation *> *commits = [[NSMutableArray alloc] init];
+  __block FIRWriteBatch *writeBatch = nil;
+  __block int writeBatchSize = 0;
+
   [documents enumerateKeysAndObjectsUsingBlock:^(NSString *key, NSDictionary<NSString *, id> *value,
                                                  BOOL *) {
-    FIRDocumentReference *ref = [collection documentWithPath:key];
-    [self writeDocumentRef:ref data:value];
+    if (writeBatch == nil) {
+      writeBatch = [collection.firestore batch];
+    }
+
+    [writeBatch setData:value forDocument:[collection documentWithPath:key]];
+    writeBatchSize++;
+
+    // Write batches are capped at 500 writes. Use 400 just to be safe.
+    if (writeBatchSize == 400) {
+      XCTestExpectation *commitExpectation = [self expectationWithDescription:@"WriteBatch commit"];
+      [writeBatch commitWithCompletion:^(NSError *_Nullable error) {
+        [commitExpectation fulfill];
+        if (error != nil) {
+          XCTFail(@"WriteBatch commit failed: %@", error);
+        }
+      }];
+      [commits addObject:commitExpectation];
+      writeBatch = nil;
+      writeBatchSize = 0;
+    }
   }];
+
+  if (writeBatch != nil) {
+    XCTestExpectation *commitExpectation = [self expectationWithDescription:@"WriteBatch commit"];
+    [writeBatch commitWithCompletion:^(NSError *_Nullable error) {
+      [commitExpectation fulfill];
+      if (error != nil) {
+        XCTFail(@"WriteBatch commit failed: %@", error);
+      }
+    }];
+    [commits addObject:commitExpectation];
+  }
+
+  for (XCTestExpectation *commitExpectation in commits) {
+    [self awaitExpectation:commitExpectation];
+  }
 }
 
 - (void)readerAndWriterOnDocumentRef:(void (^)(FIRDocumentReference *readerRef,
@@ -590,6 +628,16 @@ - (void)waitUntil:(BOOL (^)())predicate {
   return result;
 }
 
+extern "C" NSArray<FIRDocumentReference *> *FIRDocumentReferenceArrayFromQuerySnapshot(
+    FIRQuerySnapshot *docs) {
+  NSMutableArray<FIRDocumentReference *> *documentReferenceAccumulator =
+      [[NSMutableArray alloc] init];
+  for (FIRDocumentSnapshot *documentSnapshot in docs.documents) {
+    [documentReferenceAccumulator addObject:documentSnapshot.reference];
+  }
+  return [documentReferenceAccumulator copy];
+}
+
 @end
 
 NS_ASSUME_NONNULL_END