Improve bloom filter application test coverage

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

@@ -501,7 +501,6 @@ public Target encodeTarget(TargetData targetData) {
       builder.setResumeToken(targetData.getResumeToken());
     }
 
-    // TODO(Mila) Incorporate this into the if statement above.
     if (targetData.getExpectedCount() != null
         && (!targetData.getResumeToken().isEmpty()
             || targetData.getSnapshotVersion().compareTo(SnapshotVersion.NONE) > 0)) {

firebase-firestore/src/test/java/com/google/firebase/firestore/local/LocalSerializerTest.java

@@ -404,6 +404,51 @@ public void testEncodesTargetData() {
     assertEquals(targetData, decoded);
   }
 
+  @Test
+  public void localSerializerDropsExpectedCountInTargetData() {
+    Query query = TestUtil.query("room");
+    int targetId = 42;
+    long sequenceNumber = 10;
+    SnapshotVersion snapshotVersion = TestUtil.version(1039);
+    SnapshotVersion limboFreeVersion = TestUtil.version(1000);
+    ByteString resumeToken = TestUtil.resumeToken(1039);
+
+    TargetData targetData =
+        new TargetData(
+            query.toTarget(),
+            targetId,
+            sequenceNumber,
+            QueryPurpose.LISTEN,
+            snapshotVersion,
+            limboFreeVersion,
+            resumeToken,
+            /* expectedCount= */ 1234);
+
+    com.google.firestore.v1.Target.QueryTarget queryTarget =
+        remoteSerializer.encodeQueryTarget(query.toTarget());
+
+    com.google.firebase.firestore.proto.Target expected =
+        com.google.firebase.firestore.proto.Target.newBuilder()
+            .setTargetId(targetId)
+            .setLastListenSequenceNumber(sequenceNumber)
+            .setSnapshotVersion(com.google.protobuf.Timestamp.newBuilder().setNanos(1039000))
+            .setResumeToken(ByteString.copyFrom(resumeToken.toByteArray()))
+            .setQuery(
+                com.google.firestore.v1.Target.QueryTarget.newBuilder()
+                    .setParent(queryTarget.getParent())
+                    .setStructuredQuery(queryTarget.getStructuredQuery()))
+            .setLastLimboFreeSnapshotVersion(
+                com.google.protobuf.Timestamp.newBuilder().setNanos(1000000))
+            .build();
+
+    assertEquals(expected, serializer.encodeTargetData(targetData));
+    TargetData decoded = serializer.decodeTargetData(expected);
+    // Set the expected_count in TargetData to null, as serializing a TargetData into local Target
+    // proto will drop the expected_count and the deserialized TargetData will not include the
+    // expected_count.
+    assertEquals(targetData.withExpectedCount(null), decoded);
+  }
+
   @Test
   public void testEncodesQuery() {
     Target target =

firebase-firestore/src/test/java/com/google/firebase/firestore/remote/RemoteEventTest.java

@@ -39,6 +39,8 @@
 import com.google.firebase.firestore.remote.WatchChange.WatchTargetChange;
 import com.google.firebase.firestore.remote.WatchChange.WatchTargetChangeType;
 import com.google.firebase.firestore.testutil.TestTargetMetadataProvider;
+import com.google.firestore.v1.BitSequence;
+import com.google.firestore.v1.BloomFilter;
 import com.google.protobuf.ByteString;
 import java.util.ArrayList;
 import java.util.Collections;
@@ -456,6 +458,120 @@ public void testExistenceFilterMismatchClearsTarget() {
     assertEquals(0, event.getDocumentUpdates().size());
   }
 
+  @Test
+  public void existenceFilterMismatchWithBloomFilterSuccess() {
+    Map<Integer, TargetData> targetMap = activeQueries(1, 2);
+
+    MutableDocument doc1 = doc("docs/1", 1, map("value", 1));
+    MutableDocument doc2 = doc("docs/2", 2, map("value", 2));
+
+    WatchChange change1 = new DocumentChange(asList(1), emptyList(), doc1.getKey(), doc1);
+    WatchChange change2 = new DocumentChange(asList(1), emptyList(), doc2.getKey(), doc2);
+    WatchChange change3 = new WatchTargetChange(WatchTargetChangeType.Current, asList(1));
+
+    // The BloomFilter proto value below is created based on the document paths that are constructed
+    // using the pattern: "projects/test-project/databases/test-database/documents/"+document_key.
+    // Override the default database ID to ensure that the document path matches the pattern above.
+    targetMetadataProvider.setDatabaseId("test-project", "test-database");
+    WatchChangeAggregator aggregator =
+        createAggregator(
+            targetMap,
+            noOutstandingResponses,
+            keySet(doc1.getKey(), doc2.getKey()),
+            change1,
+            change2,
+            change3);
+
+    RemoteEvent event = aggregator.createRemoteEvent(version(3));
+
+    assertEquals(version(3), event.getSnapshotVersion());
+    assertEquals(2, event.getDocumentUpdates().size());
+    assertEquals(doc1, event.getDocumentUpdates().get(doc1.getKey()));
+    assertEquals(doc2, event.getDocumentUpdates().get(doc2.getKey()));
+
+    assertEquals(2, event.getTargetChanges().size());
+
+    TargetChange mapping1 = targetChange(resumeToken, true, null, asList(doc1, doc2), null);
+    assertEquals(mapping1, event.getTargetChanges().get(1));
+
+    TargetChange mapping2 = targetChange(resumeToken, false, null, null, null);
+    assertEquals(mapping2, event.getTargetChanges().get(2));
+
+    // This BloomFilter will return false on MightContain(doc1) and true on MightContain(doc2).
+    BitSequence.Builder bitSequence = BitSequence.newBuilder();
+    bitSequence.setPadding(1);
+    bitSequence.setBitmap(ByteString.copyFrom(new byte[] {0x0E, 0x0F}));
+    com.google.firestore.v1.BloomFilter.Builder bloomFilter = BloomFilter.newBuilder();
+    bloomFilter.setBits(bitSequence);
+    bloomFilter.setHashCount(7);
+
+    WatchChange.ExistenceFilterWatchChange watchChange =
+        new WatchChange.ExistenceFilterWatchChange(1, new ExistenceFilter(1, bloomFilter.build()));
+    aggregator.handleExistenceFilter(watchChange);
+
+    event = aggregator.createRemoteEvent(version(3));
+
+    assertEquals(1, event.getTargetChanges().size());
+    assertEquals(0, event.getTargetMismatches().size());
+    assertEquals(0, event.getDocumentUpdates().size());
+  }
+
+  @Test
+  public void existenceFilterMismatchWithBloomFilterFalsePositiveResult() {
+    Map<Integer, TargetData> targetMap = activeQueries(1, 2);
+
+    MutableDocument doc1 = doc("docs/1", 1, map("value", 1));
+    MutableDocument doc2 = doc("docs/2", 2, map("value", 2));
+
+    WatchChange change1 = new DocumentChange(asList(1), emptyList(), doc1.getKey(), doc1);
+    WatchChange change2 = new DocumentChange(asList(1), emptyList(), doc2.getKey(), doc2);
+    WatchChange change3 = new WatchTargetChange(WatchTargetChangeType.Current, asList(1));
+
+    WatchChangeAggregator aggregator =
+        createAggregator(
+            targetMap,
+            noOutstandingResponses,
+            keySet(doc1.getKey(), doc2.getKey()),
+            change1,
+            change2,
+            change3);
+
+    RemoteEvent event = aggregator.createRemoteEvent(version(3));
+
+    assertEquals(version(3), event.getSnapshotVersion());
+    assertEquals(2, event.getDocumentUpdates().size());
+    assertEquals(doc1, event.getDocumentUpdates().get(doc1.getKey()));
+    assertEquals(doc2, event.getDocumentUpdates().get(doc2.getKey()));
+
+    assertEquals(2, event.getTargetChanges().size());
+
+    TargetChange mapping1 = targetChange(resumeToken, true, null, asList(doc1, doc2), null);
+    assertEquals(mapping1, event.getTargetChanges().get(1));
+
+    TargetChange mapping2 = targetChange(resumeToken, false, null, null, null);
+    assertEquals(mapping2, event.getTargetChanges().get(2));
+
+    // This BloomFilter will return true on both MightContain(doc1) and MightContain(doc2).
+    BitSequence.Builder bitSequence = BitSequence.newBuilder();
+    bitSequence.setPadding(7);
+    bitSequence.setBitmap(ByteString.copyFrom(new byte[] {0x42, (byte) 0xFE}));
+    com.google.firestore.v1.BloomFilter.Builder bloomFilter = BloomFilter.newBuilder();
+    bloomFilter.setBits(bitSequence);
+    bloomFilter.setHashCount(33);
+
+    WatchChange.ExistenceFilterWatchChange watchChange =
+        new WatchChange.ExistenceFilterWatchChange(1, new ExistenceFilter(1, bloomFilter.build()));
+    aggregator.handleExistenceFilter(watchChange);
+
+    event = aggregator.createRemoteEvent(version(3));
+
+    TargetChange mapping3 = targetChange(ByteString.EMPTY, false, null, null, asList(doc1, doc2));
+    assertEquals(1, event.getTargetChanges().size());
+    assertEquals(mapping3, event.getTargetChanges().get(1));
+    assertEquals(1, event.getTargetMismatches().size());
+    assertEquals(0, event.getDocumentUpdates().size());
+  }
+
   @Test
   public void testExistenceFilterMismatchRemovesCurrentChanges() {
     Map<Integer, TargetData> targetMap = activeQueries(1);

firebase-firestore/src/test/java/com/google/firebase/firestore/remote/RemoteSerializerTest.java

@@ -509,6 +509,11 @@ public void testEncodesListenRequestLabels() {
     targetData = new TargetData(query.toTarget(), 2, 3, QueryPurpose.EXISTENCE_FILTER_MISMATCH);
     result = serializer.encodeListenRequestLabels(targetData);
     assertEquals(map("goog-listen-tags", "existence-filter-mismatch"), result);
+
+    targetData =
+        new TargetData(query.toTarget(), 2, 3, QueryPurpose.EXISTENCE_FILTER_MISMATCH_BLOOM);
+    result = serializer.encodeListenRequestLabels(targetData);
+    assertEquals(map("goog-listen-tags", "existence-filter-mismatch-bloom"), result);
   }
 
   @Test
@@ -1153,6 +1158,96 @@ public void testEncodesReadTime() {
         serializer.decodeQueryTarget(serializer.encodeQueryTarget(q.toTarget())), q.toTarget());
   }
 
+  @Test
+  public void encodesExpectedCountWhenResumeTokenIsPresent() {
+    Query q = Query.atPath(ResourcePath.fromString("docs"));
+    TargetData targetData =
+        new TargetData(q.toTarget(), 1, 2, QueryPurpose.LISTEN)
+            .withResumeToken(TestUtil.resumeToken(1000), SnapshotVersion.NONE)
+            .withExpectedCount(42);
+    Target actual = serializer.encodeTarget(targetData);
+
+    StructuredQuery.Builder structuredQueryBuilder =
+        StructuredQuery.newBuilder()
+            .addFrom(CollectionSelector.newBuilder().setCollectionId("docs"))
+            .addOrderBy(defaultKeyOrder());
+
+    QueryTarget.Builder queryBuilder =
+        QueryTarget.newBuilder()
+            .setParent("projects/p/databases/d/documents")
+            .setStructuredQuery(structuredQueryBuilder);
+    Target expected =
+        Target.newBuilder()
+            .setQuery(queryBuilder)
+            .setTargetId(1)
+            .setResumeToken(TestUtil.resumeToken(1000))
+            .setExpectedCount(Int32Value.newBuilder().setValue(42))
+            .build();
+
+    assertEquals(expected, actual);
+    assertEquals(
+        serializer.decodeQueryTarget(serializer.encodeQueryTarget(q.toTarget())), q.toTarget());
+  }
+
+  @Test
+  public void encodesExpectedCountWhenReadTimeIsPresent() {
+    Query q = Query.atPath(ResourcePath.fromString("docs"));
+    TargetData targetData =
+        new TargetData(q.toTarget(), 1, 2, QueryPurpose.LISTEN)
+            .withResumeToken(ByteString.EMPTY, version(4000000))
+            .withExpectedCount(42);
+    Target actual = serializer.encodeTarget(targetData);
+
+    StructuredQuery.Builder structuredQueryBuilder =
+        StructuredQuery.newBuilder()
+            .addFrom(CollectionSelector.newBuilder().setCollectionId("docs"))
+            .addOrderBy(defaultKeyOrder());
+
+    QueryTarget.Builder queryBuilder =
+        QueryTarget.newBuilder()
+            .setParent("projects/p/databases/d/documents")
+            .setStructuredQuery(structuredQueryBuilder);
+    Target expected =
+        Target.newBuilder()
+            .setQuery(queryBuilder)
+            .setTargetId(1)
+            .setReadTime(Timestamp.newBuilder().setSeconds(4))
+            .setExpectedCount(Int32Value.newBuilder().setValue(42))
+            .build();
+
+    assertEquals(expected, actual);
+    assertEquals(
+        serializer.decodeQueryTarget(serializer.encodeQueryTarget(q.toTarget())), q.toTarget());
+  }
+
+  @Test
+  public void ignoresExpectedCountWithoutResumeTokenOrReadTime() {
+    Query q = Query.atPath(ResourcePath.fromString("docs"));
+    TargetData targetData =
+        new TargetData(q.toTarget(), 1, 2, QueryPurpose.LISTEN).withExpectedCount(42);
+    Target actual = serializer.encodeTarget(targetData);
+
+    StructuredQuery.Builder structuredQueryBuilder =
+        StructuredQuery.newBuilder()
+            .addFrom(CollectionSelector.newBuilder().setCollectionId("docs"))
+            .addOrderBy(defaultKeyOrder());
+
+    QueryTarget.Builder queryBuilder =
+        QueryTarget.newBuilder()
+            .setParent("projects/p/databases/d/documents")
+            .setStructuredQuery(structuredQueryBuilder);
+    Target expected =
+        Target.newBuilder()
+            .setQuery(queryBuilder)
+            .setTargetId(1)
+            .setResumeToken(ByteString.EMPTY)
+            .build();
+
+    assertEquals(expected, actual);
+    assertEquals(
+        serializer.decodeQueryTarget(serializer.encodeQueryTarget(q.toTarget())), q.toTarget());
+  }
+
   /**
    * Wraps the given query in TargetData. This is useful because the APIs we're testing accept
    * TargetData, but for the most part we're just testing variations on Query.

firebase-firestore/src/testUtil/java/com/google/firebase/firestore/testutil/TestTargetMetadataProvider.java

@@ -30,6 +30,7 @@
 public class TestTargetMetadataProvider implements WatchChangeAggregator.TargetMetadataProvider {
   final Map<Integer, ImmutableSortedSet<DocumentKey>> syncedKeys = new HashMap<>();
   final Map<Integer, TargetData> queryData = new HashMap<>();
+  DatabaseId databaseId = DatabaseId.forProject("test-project");
 
   @Override
   public ImmutableSortedSet<DocumentKey> getRemoteKeysForTarget(int targetId) {
@@ -44,7 +45,12 @@ public TargetData getTargetDataForTarget(int targetId) {
 
   @Override
   public DatabaseId getDatabaseId() {
-    return DatabaseId.forProject("test-project");
+    return databaseId;
+  }
+
+  /** Replaces the default project ID and database ID. */
+  public void setDatabaseId(String projectId, String databaseId) {
+    this.databaseId = DatabaseId.forDatabase(projectId, databaseId);
   }
 
   /** Sets or replaces the local state for the provided query data. */