Cleanup and improvements to the "complex" bloom filter integration test

packages/firestore/test/integration/api/query.test.ts

@@ -2065,115 +2065,137 @@ apiDescribe('Queries', (persistence: boolean) => {
   });
 
   it('resuming a query should use bloom filter to avoid full requery', async () => {
-    // Create 100 documents in a new collection.
+    // Prepare the names and contents of the 100 documents to create.
     const testDocs: { [key: string]: object } = {};
-    for (let i = 1; i <= 100; i++) {
-      testDocs['doc' + i] = { key: i };
+    for (let i = 0; i < 100; i++) {
+      testDocs['doc' + (1000 + i)] = { key: 42 };
     }
 
     // The function that runs a single iteration of the test.
-    // Below this definition, there is a "while" loop that calls this
-    // function potentially multiple times.
+    // Below this definition, there is a "while" loop that calls this function
+    // potentially multiple times.
     const runTestIteration = async (
       coll: CollectionReference,
       db: Firestore
     ): Promise<'retry' | 'passed'> => {
-      // Run a query to populate the local cache with the 100 documents
-      // and a resume token.
+      // Run a query to populate the local cache with the 100 documents and a
+      // resume token.
       const snapshot1 = await getDocs(coll);
       expect(snapshot1.size, 'snapshot1.size').to.equal(100);
+      const createdDocuments = snapshot1.docs.map(snapshot => snapshot.ref);
 
-      // Delete 50 of the 100 documents. Do this in a transaction, rather
-      // than deleteDoc(), to avoid affecting the local cache.
+      // Delete 50 of the 100 documents. Do this in a transaction, rather than
+      // deleteDoc(), to avoid affecting the local cache.
+      const deletedDocumentIds = new Set<string>();
       await runTransaction(db, async txn => {
-        for (let i = 1; i <= 50; i++) {
-          txn.delete(doc(coll, 'doc' + i));
+        for (let i = 0; i < createdDocuments.length; i += 2) {
+          const documentToDelete = createdDocuments[i];
+          txn.delete(documentToDelete);
+          deletedDocumentIds.add(documentToDelete.id);
         }
       });
 
-      // 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.
+      // 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.
       await new Promise(resolve => setTimeout(resolve, 10000));
 
-      // Resume the query and expect to get a snapshot with the 50
-      // remaining documents. Use some internal testing hooks to "capture"
-      // the existence filter mismatches to later verify that Watch sent a
-      // bloom filter, and it was used to avert a full requery.
-      const existenceFilterMismatches = await captureExistenceFilterMismatches(
-        async () => {
-          const snapshot2 = await getDocs(coll);
-          // 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 (!(USE_EMULATOR && snapshot2.size === 100)) {
-            expect(snapshot2.size, 'snapshot2.size').to.equal(50);
-          }
-        }
-      );
+      // Resume the query and save the resulting snapshot for verification.
+      // Use some internal testing hooks to "capture" the existence filter
+      // mismatches to verify that Watch sent a bloom filter, and it was used to
+      // avert a full requery.
+      const [existenceFilterMismatches, snapshot2] =
+        await captureExistenceFilterMismatches(() => getDocs(coll));
+
+      // 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 (!(USE_EMULATOR && snapshot2.size === 100)) {
+        const actualDocumentIds = snapshot2.docs
+          .map(documentSnapshot => documentSnapshot.ref.id)
+          .sort();
+        const expectedDocumentIds = createdDocuments
+          .filter(documentRef => !deletedDocumentIds.has(documentRef.id))
+          .map(documentRef => documentRef.id)
+          .sort();
+        expect(actualDocumentIds, 'snapshot2.docs').to.deep.equal(
+          expectedDocumentIds
+        );
+      }
 
-      // Skip the verification of the existence filter mismatch when
-      // persistence is disabled because without persistence there is no
-      // resume token specified in the subsequent call to getDocs(), and,
-      // therefore, Watch will _not_ send an existence filter.
+      // Skip the verification of the existence filter mismatch when persistence
+      // is disabled because without persistence there is no resume token
+      // specified in the subsequent call to getDocs(), and, therefore, Watch
+      // will _not_ send an existence filter.
+      // TODO(b/272754156) Re-write this test using a snapshot listener instead
+      // of calls to getDocs() and remove this check for disabled persistence.
       if (!persistence) {
         return 'passed';
       }
 
-      // Skip the verification of the existence filter mismatch when
-      // testing against the Firestore emulator because the Firestore
-      // emulator does not include the `unchanged_names` bloom filter when
-      // it sends ExistenceFilter messages. Some day the emulator _may_
-      // implement this logic, at which time this short-circuit can be
-      // removed.
+      // Skip the verification of the existence filter mismatch when testing
+      // against the Firestore emulator because the Firestore emulator does not
+      // include the `unchanged_names` bloom filter when it sends
+      // ExistenceFilter messages. Some day the emulator _may_ implement this
+      // logic, at which time this short-circuit can be removed.
       if (USE_EMULATOR) {
         return 'passed';
       }
 
-      // Verify that upon resuming the query that Watch sent an existence
-      // filter that included a bloom filter, and that the bloom filter
-      // was successfully used to avoid a full requery.
-      // TODO(b/271949433) Remove this check for "nightly" once the bloom
-      // filter logic is deployed to production, circa May 2023.
-      if (TARGET_BACKEND === 'nightly') {
-        expect(
-          existenceFilterMismatches,
-          'existenceFilterMismatches'
-        ).to.have.length(1);
-        const { localCacheCount, existenceFilterCount, bloomFilter } =
-          existenceFilterMismatches[0];
-
-        expect(localCacheCount, 'localCacheCount').to.equal(100);
-        expect(existenceFilterCount, 'existenceFilterCount').to.equal(50);
-        if (!bloomFilter) {
-          expect.fail(
-            'The existence filter should have specified ' +
-              'a bloom filter in its `unchanged_names` field.'
-          );
-          throw new Error('should never get here');
-        }
+      // Verify that Watch sent an existence filter with the correct counts when
+      // the query was resumed.
+      expect(
+        existenceFilterMismatches,
+        'existenceFilterMismatches'
+      ).to.have.length(1);
+      const { localCacheCount, existenceFilterCount, bloomFilter } =
+        existenceFilterMismatches[0];
+      expect(localCacheCount, 'localCacheCount').to.equal(100);
+      expect(existenceFilterCount, 'existenceFilterCount').to.equal(50);
+
+      // Skip the verification of the bloom filter when testing against
+      // production because the bloom filter is only implemented in nightly.
+      // TODO(b/271949433) Remove this "if" block once the bloom filter logic is
+      // deployed to production.
+      if (TARGET_BACKEND !== 'nightly') {
+        return 'passed';
+      }
 
-        expect(bloomFilter.hashCount, 'bloomFilter.hashCount').to.be.above(0);
-        expect(
-          bloomFilter.bitmapLength,
-          'bloomFilter.bitmapLength'
-        ).to.be.above(0);
-        expect(bloomFilter.padding, 'bloomFilterPadding').to.be.above(0);
-        expect(bloomFilter.padding, 'bloomFilterPadding').to.be.below(8);
-
-        // Retry the entire test if a bloom filter false positive occurs.
-        // Although statistically rare, false positives are expected to
-        // happen occasionally. When a false positive _does_ happen, just
-        // retry the test with a different set of documents. If that retry
-        // _also_ experiences a false positive, then fail the test because
-        // that is so improbable that something must have gone wrong.
-        if (attemptNumber > 1 && !bloomFilter.applied) {
-          return 'retry';
-        }
-        expect(bloomFilter.applied, 'bloomFilter.applied').to.be.true;
+      // Verify that Watch sent a valid bloom filter.
+      if (!bloomFilter) {
+        expect.fail(
+          'The existence filter should have specified a bloom filter in its ' +
+            '`unchanged_names` field.'
+        );
+        throw new Error('should never get here');
+      }
+
+      expect(bloomFilter.hashCount, 'bloomFilter.hashCount').to.be.above(0);
+      expect(bloomFilter.bitmapLength, 'bloomFilter.bitmapLength').to.be.above(
+        0
+      );
+      expect(bloomFilter.padding, 'bloomFilterPadding').to.be.above(0);
+      expect(bloomFilter.padding, 'bloomFilterPadding').to.be.below(8);
+
+      // Verify that the bloom filter was successfully used to avert a full
+      // requery. If a false positive occurred then retry the entire test.
+      // Although statistically rare, false positives are expected to happen
+      // occasionally. When a false positive _does_ happen, just retry the test
+      // with a different set of documents. If that retry _also_ experiences a
+      // false positive, then fail the test because that is so improbable that
+      // something must have gone wrong.
+      if (attemptNumber === 1 && !bloomFilter.applied) {
+        return 'retry';
       }
+      expect(
+        bloomFilter.applied,
+        `bloomFilter.applied with attemptNumber=${attemptNumber}`
+      ).to.be.true;
 
       return 'passed';
     };

packages/firestore/test/integration/util/helpers.ts

@@ -32,7 +32,9 @@ import {
   PrivateSettings,
   SnapshotListenOptions,
   newTestFirestore,
-  newTestApp
+  newTestApp,
+  writeBatch,
+  WriteBatch
 } from './firebase_export';
 import {
   ALT_PROJECT_ID,
@@ -315,11 +317,34 @@ export function withTestCollectionSettings<T>(
       const collectionId = 'test-collection-' + doc(collection(testDb, 'x')).id;
       const testCollection = collection(testDb, collectionId);
       const setupCollection = collection(setupDb, collectionId);
-      const sets: Array<Promise<void>> = [];
-      Object.keys(docs).forEach(key => {
-        sets.push(setDoc(doc(setupCollection, key), docs[key]));
-      });
-      return Promise.all(sets).then(() => fn(testCollection, testDb));
+
+      const writeBatchCommits: Array<Promise<void>> = [];
+      let writeBatch_: WriteBatch | null = null;
+      let writeBatchSize = 0;
+
+      for (const key of Object.keys(docs)) {
+        if (writeBatch_ === null) {
+          writeBatch_ = writeBatch(setupDb);
+        }
+
+        writeBatch_.set(doc(setupCollection, key), docs[key]);
+        writeBatchSize++;
+
+        // Write batches are capped at 500 writes. Use 400 just to be safe.
+        if (writeBatchSize === 400) {
+          writeBatchCommits.push(writeBatch_.commit());
+          writeBatch_ = null;
+          writeBatchSize = 0;
+        }
+      }
+
+      if (writeBatch_ !== null) {
+        writeBatchCommits.push(writeBatch_.commit());
+      }
+
+      return Promise.all(writeBatchCommits).then(() =>
+        fn(testCollection, testDb)
+      );
     }
   );
 }

packages/firestore/test/integration/util/testing_hooks_util.ts

@@ -24,9 +24,9 @@ import { _TestingHooks as TestingHooks } from './firebase_export';
  * callback all existence filter mismatches will be captured.
  * @return the captured existence filter mismatches.
  */
-export async function captureExistenceFilterMismatches(
-  callback: () => Promise<void>
-): Promise<ExistenceFilterMismatchInfo[]> {
+export async function captureExistenceFilterMismatches<T>(
+  callback: () => Promise<T>
+): Promise<[ExistenceFilterMismatchInfo[], T]> {
   const results: ExistenceFilterMismatchInfo[] = [];
   const onExistenceFilterMismatchCallback = (
     info: ExistenceFilterMismatchInfo
@@ -39,13 +39,14 @@ export async function captureExistenceFilterMismatches(
       onExistenceFilterMismatchCallback
     );
 
+  let callbackResult: T;
   try {
-    await callback();
+    callbackResult = await callback();
   } finally {
     unregister();
   }
 
-  return results;
+  return [results, callbackResult];
 }
 
 /**