Firestore: Update test bloomFilterShouldCorrectlyEncodeComplexUnicodeCharacters() in QueryTest.java

firebase-firestore/src/androidTest/java/com/google/firebase/firestore/QueryTest.java

@@ -1220,109 +1220,104 @@ public void bloomFilterShouldCorrectlyEncodeComplexUnicodeCharacters() throws Ex
       testDocs.put(docId, map("foo", 42));
     }
 
-    // Each iteration of the "while" loop below runs a single iteration of the test. The test will
-    // be run multiple times only if a bloom filter false positive occurs.
-    int attemptNumber = 0;
-    while (true) {
-      attemptNumber++;
-
-      // Create the documents whose names contain complex Unicode characters in a new collection.
-      CollectionReference collection = testCollectionWithDocs(testDocs);
-
-      // Run a query to populate the local cache with documents that have names with complex Unicode
-      // characters.
-      List<DocumentReference> createdDocuments = new ArrayList<>();
-      {
-        QuerySnapshot querySnapshot1 = waitFor(collection.get());
-        for (DocumentSnapshot documentSnapshot : querySnapshot1.getDocuments()) {
-          createdDocuments.add(documentSnapshot.getReference());
-        }
-        HashSet<String> createdDocumentIds = new HashSet<>();
-        for (DocumentSnapshot documentSnapshot : querySnapshot1.getDocuments()) {
-          createdDocumentIds.add(documentSnapshot.getId());
-        }
-        assertWithMessage("createdDocumentIds")
-            .that(createdDocumentIds)
-            .containsExactlyElementsIn(testDocIds);
-      }
-
-      // Delete one of the documents so that the next call to getDocs() will
-      // experience an existence filter mismatch. Do this deletion in a
-      // transaction, rather than using deleteDoc(), to avoid affecting the
-      // local cache.
-      waitFor(
-          collection
-              .getFirestore()
-              .runTransaction(
-                  transaction -> {
-                    DocumentReference documentToDelete = collection.document("DocumentToDelete");
-                    DocumentSnapshot documentToDeleteSnapshot = transaction.get(documentToDelete);
-                    assertWithMessage("documentToDeleteSnapshot.exists()")
-                        .that(documentToDeleteSnapshot.exists())
-                        .isTrue();
-                    transaction.delete(documentToDelete);
-                    return null;
-                  }));
-
-      // 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.
-      Thread.sleep(10000);
-
-      // Resume the query and save the resulting snapshot for verification. Use some internal
-      // testing hooks to "capture" the existence filter mismatches.
-      AtomicReference<QuerySnapshot> querySnapshot2Ref = new AtomicReference<>();
-      ArrayList<ExistenceFilterMismatchInfo> existenceFilterMismatches =
-          captureExistenceFilterMismatches(
-              () -> {
-                QuerySnapshot querySnapshot = waitFor(collection.get());
-                querySnapshot2Ref.set(querySnapshot);
-              });
-      QuerySnapshot querySnapshot2 = querySnapshot2Ref.get();
+    // Create the documents whose names contain complex Unicode characters in a new collection.
+    CollectionReference collection = testCollectionWithDocs(testDocs);
 
-      // Verify that the snapshot from the resumed query contains the expected documents; that is,
-      // that it contains the documents whose names contain complex Unicode characters and _not_ the
-      // document that was deleted.
-      HashSet<String> querySnapshot2DocumentIds = new HashSet<>();
-      for (DocumentSnapshot documentSnapshot : querySnapshot2.getDocuments()) {
-        querySnapshot2DocumentIds.add(documentSnapshot.getId());
+    // Run a query to populate the local cache with documents that have names with complex Unicode
+    // characters.
+    List<DocumentReference> createdDocuments = new ArrayList<>();
+    {
+      QuerySnapshot querySnapshot1 = waitFor(collection.get());
+      for (DocumentSnapshot documentSnapshot : querySnapshot1.getDocuments()) {
+        createdDocuments.add(documentSnapshot.getReference());
       }
-      HashSet<String> querySnapshot2ExpectedDocumentIds = new HashSet<>(testDocIds);
-      querySnapshot2ExpectedDocumentIds.remove("DocumentToDelete");
-      assertWithMessage("querySnapshot2DocumentIds")
-          .that(querySnapshot2DocumentIds)
-          .containsExactlyElementsIn(querySnapshot2ExpectedDocumentIds);
-
-      // Verify that Watch sent an existence filter with the correct counts.
-      assertWithMessage("Watch should have sent exactly 1 existence filter")
-          .that(existenceFilterMismatches)
-          .hasSize(1);
-      ExistenceFilterMismatchInfo existenceFilterMismatchInfo = existenceFilterMismatches.get(0);
-      assertWithMessage("localCacheCount")
-          .that(existenceFilterMismatchInfo.localCacheCount())
-          .isEqualTo(testDocIds.size());
-      assertWithMessage("existenceFilterCount")
-          .that(existenceFilterMismatchInfo.existenceFilterCount())
-          .isEqualTo(testDocIds.size() - 1);
-
-      // Verify that Watch sent a valid bloom filter.
-      ExistenceFilterBloomFilterInfo bloomFilter = existenceFilterMismatchInfo.bloomFilter();
-      assertWithMessage("The bloom filter specified in the existence filter")
-          .that(bloomFilter)
-          .isNotNull();
-
-      // Verify that the bloom filter was successfully used to avert a full requery. If a false
-      // positive occurred, which is statistically rare, but technically possible, then retry the
-      // entire test.
-      if (attemptNumber == 1 && !bloomFilter.applied()) {
-        continue;
+      HashSet<String> createdDocumentIds = new HashSet<>();
+      for (DocumentSnapshot documentSnapshot : querySnapshot1.getDocuments()) {
+        createdDocumentIds.add(documentSnapshot.getId());
       }
+      assertWithMessage("createdDocumentIds")
+          .that(createdDocumentIds)
+          .containsExactlyElementsIn(testDocIds);
+    }
 
-      assertWithMessage("bloom filter successfully applied with attemptNumber=" + attemptNumber)
-          .that(bloomFilter.applied())
+    // Delete one of the documents so that the next call to getDocs() will experience an existence
+    // filter mismatch. Do this deletion in a transaction, rather than using deleteDoc(), to avoid
+    // affecting the local cache.
+    DocumentReference documentToDelete = collection.document("DocumentToDelete");
+    waitFor(
+        collection
+            .getFirestore()
+            .runTransaction(
+                transaction -> {
+                  DocumentSnapshot documentToDeleteSnapshot = transaction.get(documentToDelete);
+                  assertWithMessage("documentToDeleteSnapshot.exists()")
+                      .that(documentToDeleteSnapshot.exists())
+                      .isTrue();
+                  transaction.delete(documentToDelete);
+                  return null;
+                }));
+
+    // 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.
+    Thread.sleep(10000);
+
+    // Resume the query and save the resulting snapshot for verification. Use some internal testing
+    // hooks to "capture" the existence filter mismatches.
+    AtomicReference<QuerySnapshot> querySnapshot2Ref = new AtomicReference<>();
+    ArrayList<ExistenceFilterMismatchInfo> existenceFilterMismatches =
+        captureExistenceFilterMismatches(
+            () -> {
+              QuerySnapshot querySnapshot = waitFor(collection.get());
+              querySnapshot2Ref.set(querySnapshot);
+            });
+    QuerySnapshot querySnapshot2 = querySnapshot2Ref.get();
+
+    // Verify that the snapshot from the resumed query contains the expected documents; that is,
+    // that it contains the documents whose names contain complex Unicode characters and _not_ the
+    // document that was deleted.
+    HashSet<String> querySnapshot2DocumentIds = new HashSet<>();
+    for (DocumentSnapshot documentSnapshot : querySnapshot2.getDocuments()) {
+      querySnapshot2DocumentIds.add(documentSnapshot.getId());
+    }
+    HashSet<String> querySnapshot2ExpectedDocumentIds = new HashSet<>(testDocIds);
+    querySnapshot2ExpectedDocumentIds.remove("DocumentToDelete");
+    assertWithMessage("querySnapshot2DocumentIds")
+        .that(querySnapshot2DocumentIds)
+        .containsExactlyElementsIn(querySnapshot2ExpectedDocumentIds);
+
+    // Verify that Watch sent an existence filter with the correct counts.
+    assertWithMessage("Watch should have sent exactly 1 existence filter")
+        .that(existenceFilterMismatches)
+        .hasSize(1);
+    ExistenceFilterMismatchInfo existenceFilterMismatchInfo = existenceFilterMismatches.get(0);
+    assertWithMessage("localCacheCount")
+        .that(existenceFilterMismatchInfo.localCacheCount())
+        .isEqualTo(testDocIds.size());
+    assertWithMessage("existenceFilterCount")
+        .that(existenceFilterMismatchInfo.existenceFilterCount())
+        .isEqualTo(testDocIds.size() - 1);
+
+    // Verify that Watch sent a valid bloom filter.
+    ExistenceFilterBloomFilterInfo bloomFilter = existenceFilterMismatchInfo.bloomFilter();
+    assertWithMessage("The bloom filter specified in the existence filter")
+        .that(bloomFilter)
+        .isNotNull();
+
+    // The bloom filter application should statistically be successful almost every time; the _only_
+    // time when it would _not_ be successful is if there is a false positive when testing for
+    // 'DocumentToDelete' in the bloom filter. So verify that the bloom filter application is
+    // successful, unless there was a false positive.
+    boolean isFalsePositive = bloomFilter.mightContain(documentToDelete);
+    assertWithMessage("bloomFilter.applied()")
+        .that(bloomFilter.applied())
+        .isEqualTo(!isFalsePositive);
+
+    // Verify that the bloom filter contains the document paths with complex Unicode characters.
+    for (DocumentSnapshot documentSnapshot : querySnapshot2.getDocuments()) {
+      DocumentReference documentReference = documentSnapshot.getReference();
+      assertWithMessage("bloomFilter.mightContain() for " + documentReference.getPath())
+          .that(bloomFilter.mightContain(documentReference))
           .isTrue();
-
-      // Break out of the test loop now that the test passes.
-      break;
     }
   }
 

firebase-firestore/src/androidTest/java/com/google/firebase/firestore/remote/TestingHooksUtil.java

@@ -16,6 +16,7 @@
 
 import androidx.annotation.NonNull;
 import androidx.annotation.Nullable;
+import com.google.firebase.firestore.DocumentReference;
 import com.google.firebase.firestore.ListenerRegistration;
 import java.util.ArrayList;
 
@@ -84,17 +85,27 @@ public int existenceFilterCount() {
     @Nullable
     public ExistenceFilterBloomFilterInfo bloomFilter() {
       TestingHooks.ExistenceFilterBloomFilterInfo bloomFilterInfo = info.bloomFilter();
-      return bloomFilterInfo == null ? null : new ExistenceFilterBloomFilterInfo(bloomFilterInfo);
+      return bloomFilterInfo == null
+          ? null
+          : new ExistenceFilterBloomFilterInfo(
+              bloomFilterInfo, info.projectId(), info.databaseId());
     }
   }
 
   /** @see TestingHooks.ExistenceFilterBloomFilterInfo */
   public static final class ExistenceFilterBloomFilterInfo {
 
     private final TestingHooks.ExistenceFilterBloomFilterInfo info;
+    private final String projectId;
+    private final String databaseId;
 
-    ExistenceFilterBloomFilterInfo(@NonNull TestingHooks.ExistenceFilterBloomFilterInfo info) {
+    ExistenceFilterBloomFilterInfo(
+        @NonNull TestingHooks.ExistenceFilterBloomFilterInfo info,
+        String projectId,
+        String databaseId) {
       this.info = info;
+      this.projectId = projectId;
+      this.databaseId = databaseId;
     }
 
     public boolean applied() {
@@ -112,5 +123,22 @@ public int bitmapLength() {
     public int padding() {
       return info.padding();
     }
+
+    public boolean mightContain(DocumentReference documentReference) {
+      BloomFilter bloomFilter = info.bloomFilter();
+      if (bloomFilter == null) {
+        return false;
+      }
+      return bloomFilter.mightContain(getBloomFilterEntryFor(documentReference));
+    }
+
+    private String getBloomFilterEntryFor(DocumentReference documentReference) {
+      return "projects/"
+          + projectId
+          + "/databases/"
+          + databaseId
+          + "/documents/"
+          + documentReference.getPath();
+    }
   }
 }

firebase-firestore/src/main/java/com/google/firebase/firestore/remote/BloomFilter.java

@@ -16,7 +16,6 @@
 
 import android.util.Base64;
 import androidx.annotation.NonNull;
-import androidx.annotation.VisibleForTesting;
 import com.google.protobuf.ByteString;
 import java.nio.charset.StandardCharsets;
 import java.security.MessageDigest;
@@ -93,7 +92,6 @@ public BloomFilterCreateException(String message) {
     }
   }
 
-  @VisibleForTesting
   int getBitCount() {
     return this.bitCount;
   }

firebase-firestore/src/main/java/com/google/firebase/firestore/remote/TestingHooks.java

@@ -22,6 +22,7 @@
 import androidx.annotation.VisibleForTesting;
 import com.google.auto.value.AutoValue;
 import com.google.firebase.firestore.ListenerRegistration;
+import com.google.firebase.firestore.model.DatabaseId;
 import com.google.firestore.v1.BloomFilter;
 import java.util.concurrent.CopyOnWriteArrayList;
 import java.util.concurrent.atomic.AtomicReference;
@@ -128,9 +129,11 @@ abstract static class ExistenceFilterMismatchInfo {
     static ExistenceFilterMismatchInfo create(
         int localCacheCount,
         int existenceFilterCount,
+        String projectId,
+        String databaseId,
         @Nullable ExistenceFilterBloomFilterInfo bloomFilter) {
       return new AutoValue_TestingHooks_ExistenceFilterMismatchInfo(
-          localCacheCount, existenceFilterCount, bloomFilter);
+          localCacheCount, existenceFilterCount, projectId, databaseId, bloomFilter);
     }
 
     /** Returns the number of documents that matched the query in the local cache. */
@@ -142,6 +145,12 @@ static ExistenceFilterMismatchInfo create(
      */
     abstract int existenceFilterCount();
 
+    /** The projectId used when checking documents for membership in the bloom filter. */
+    abstract String projectId();
+
+    /** The databaseId used when checking documents for membership in the bloom filter. */
+    abstract String databaseId();
+
     /**
      * Returns information about the bloom filter provided by Watch in the ExistenceFilter message's
      * `unchangedNames` field. A `null` return value means that Watch did _not_ provide a bloom
@@ -155,23 +164,37 @@ static ExistenceFilterMismatchInfo create(
      * with the values taken from the given arguments.
      */
     static ExistenceFilterMismatchInfo from(
-        boolean bloomFilterApplied, int localCacheCount, ExistenceFilter existenceFilter) {
+        int localCacheCount,
+        ExistenceFilter existenceFilter,
+        DatabaseId databaseId,
+        @Nullable com.google.firebase.firestore.remote.BloomFilter bloomFilter,
+        boolean bloomFilterApplied) {
       return create(
           localCacheCount,
           existenceFilter.getCount(),
-          ExistenceFilterBloomFilterInfo.from(bloomFilterApplied, existenceFilter));
+          databaseId.getProjectId(),
+          databaseId.getDatabaseId(),
+          ExistenceFilterBloomFilterInfo.from(bloomFilter, bloomFilterApplied, existenceFilter));
     }
   }
 
   @AutoValue
   abstract static class ExistenceFilterBloomFilterInfo {
 
     static ExistenceFilterBloomFilterInfo create(
-        boolean applied, int hashCount, int bitmapLength, int padding) {
+        @Nullable com.google.firebase.firestore.remote.BloomFilter bloomFilter,
+        boolean applied,
+        int hashCount,
+        int bitmapLength,
+        int padding) {
       return new AutoValue_TestingHooks_ExistenceFilterBloomFilterInfo(
-          applied, hashCount, bitmapLength, padding);
+          bloomFilter, applied, hashCount, bitmapLength, padding);
     }
 
+    /** The BloomFilter created from the existence filter; may be null if creating it failed. */
+    @Nullable
+    abstract com.google.firebase.firestore.remote.BloomFilter bloomFilter();
+
     /**
      * Returns whether a full requery was averted by using the bloom filter. If false, then
      * something happened, such as a false positive, to prevent using the bloom filter to avoid a
@@ -188,13 +211,17 @@ static ExistenceFilterBloomFilterInfo create(
     /** Returns the number of bits of padding in the last byte of the bloom filter. */
     abstract int padding();
 
+    @Nullable
     static ExistenceFilterBloomFilterInfo from(
-        boolean bloomFilterApplied, ExistenceFilter existenceFilter) {
+        @Nullable com.google.firebase.firestore.remote.BloomFilter bloomFilter,
+        boolean bloomFilterApplied,
+        ExistenceFilter existenceFilter) {
       BloomFilter unchangedNames = existenceFilter.getUnchangedNames();
       if (unchangedNames == null) {
         return null;
       }
       return create(
+          bloomFilter,
           bloomFilterApplied,
           unchangedNames.getHashCount(),
           unchangedNames.getBits().getBitmap().size(),