Improve the integration test coverage for online vs offline comparisons.

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

@@ -91,24 +91,27 @@ public void testOrQueriesWithCompositeIndexes() {
 
     Query query = collection.where(or(greaterThan("a", 2), equalTo("b", 1)));
     // with one inequality: a>2 || b==1.
-    testHelper.assertOnlineAndOfflineResultsMatch(testHelper.query(query), "doc5", "doc2", "doc3");
+    testHelper.assertOnlineAndOfflineResultsMatch(
+        collection, testHelper.query(query), "doc5", "doc2", "doc3");
 
     // Test with limits (implicit order by ASC): (a==1) || (b > 0) LIMIT 2
     query = collection.where(or(equalTo("a", 1), greaterThan("b", 0))).limit(2);
-    testHelper.assertOnlineAndOfflineResultsMatch(testHelper.query(query), "doc1", "doc2");
+    testHelper.assertOnlineAndOfflineResultsMatch(
+        collection, testHelper.query(query), "doc1", "doc2");
 
     // Test with limits (explicit order by): (a==1) || (b > 0) LIMIT_TO_LAST 2
     // Note: The public query API does not allow implicit ordering when limitToLast is used.
     query = collection.where(or(equalTo("a", 1), greaterThan("b", 0))).limitToLast(2).orderBy("b");
-    testHelper.assertOnlineAndOfflineResultsMatch(testHelper.query(query), "doc3", "doc4");
+    testHelper.assertOnlineAndOfflineResultsMatch(
+        collection, testHelper.query(query), "doc3", "doc4");
 
     // Test with limits (explicit order by ASC): (a==2) || (b == 1) ORDER BY a LIMIT 1
     query = collection.where(or(equalTo("a", 2), equalTo("b", 1))).limit(1).orderBy("a");
-    testHelper.assertOnlineAndOfflineResultsMatch(testHelper.query(query), "doc5");
+    testHelper.assertOnlineAndOfflineResultsMatch(collection, testHelper.query(query), "doc5");
 
     // Test with limits (explicit order by DESC): (a==2) || (b == 1) ORDER BY a LIMIT_TO_LAST 1
     query = collection.where(or(equalTo("a", 2), equalTo("b", 1))).limitToLast(1).orderBy("a");
-    testHelper.assertOnlineAndOfflineResultsMatch(testHelper.query(query), "doc2");
+    testHelper.assertOnlineAndOfflineResultsMatch(collection, testHelper.query(query), "doc2");
   }
 
   @Test
@@ -771,17 +774,17 @@ public void testMultipleInequalityFromCacheAndFromServer() {
 
     // implicit AND: a != 1 && b < 2
     Query query1 = testHelper.query(collection).whereNotEqualTo("a", 1).whereLessThan("b", 2);
-    testHelper.assertOnlineAndOfflineResultsMatch(query1, "doc2");
+    testHelper.assertOnlineAndOfflineResultsMatch(collection, query1, "doc2");
 
     // explicit AND: a != 1 && b < 2
     Query query2 = testHelper.query(collection).where(and(notEqualTo("a", 1), lessThan("b", 2)));
-    testHelper.assertOnlineAndOfflineResultsMatch(query2, "doc2");
+    testHelper.assertOnlineAndOfflineResultsMatch(collection, query2, "doc2");
 
     // explicit AND: a < 3 && b not-in [2, 3]
     // Implicitly ordered by: a asc, b asc, __name__ asc
     Query query3 =
         testHelper.query(collection).where(and(lessThan("a", 3), notInArray("b", asList(2, 3))));
-    testHelper.assertOnlineAndOfflineResultsMatch(query3, "doc1", "doc5", "doc2");
+    testHelper.assertOnlineAndOfflineResultsMatch(collection, query3, "doc1", "doc5", "doc2");
 
     // a <3 && b != 0, ordered by: b desc, a desc, __name__ desc
     Query query4 =
@@ -791,11 +794,11 @@ public void testMultipleInequalityFromCacheAndFromServer() {
             .whereNotEqualTo("b", 0)
             .orderBy("b", Direction.DESCENDING)
             .limit(2);
-    testHelper.assertOnlineAndOfflineResultsMatch(query4, "doc4", "doc2");
+    testHelper.assertOnlineAndOfflineResultsMatch(collection, query4, "doc4", "doc2");
 
     // explicit OR: a>2 || b<1.
     Query query5 = testHelper.query(collection).where(or(greaterThan("a", 2), lessThan("b", 1)));
-    testHelper.assertOnlineAndOfflineResultsMatch(query5, "doc1", "doc3");
+    testHelper.assertOnlineAndOfflineResultsMatch(collection, query5, "doc1", "doc3");
   }
 
   @Test

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

@@ -1651,7 +1651,7 @@ public void sdkOrdersQueryByDocumentIdTheSameWayOnlineAndOffline() {
             "a");
 
     // Run query with snapshot listener
-    checkOnlineAndOfflineResultsMatch(orderedQuery, expectedDocIds.toArray(new String[0]));
+    checkOnlineAndOfflineResultsMatch(colRef, orderedQuery, expectedDocIds.toArray(new String[0]));
   }
 
   @Test
@@ -1708,7 +1708,7 @@ public void snapshotListenerSortsUnicodeStringsAsServer() {
     assertTrue(getSnapshotDocIds.equals(expectedDocIds));
     assertTrue(watchSnapshotDocIds.equals(expectedDocIds));
 
-    checkOnlineAndOfflineResultsMatch(orderedQuery, expectedDocIds.toArray(new String[0]));
+    checkOnlineAndOfflineResultsMatch(colRef, orderedQuery, expectedDocIds.toArray(new String[0]));
   }
 
   @Test
@@ -1765,7 +1765,7 @@ public void snapshotListenerSortsUnicodeStringsInArrayAsServer() {
     assertTrue(getSnapshotDocIds.equals(expectedDocIds));
     assertTrue(watchSnapshotDocIds.equals(expectedDocIds));
 
-    checkOnlineAndOfflineResultsMatch(orderedQuery, expectedDocIds.toArray(new String[0]));
+    checkOnlineAndOfflineResultsMatch(colRef, orderedQuery, expectedDocIds.toArray(new String[0]));
   }
 
   @Test
@@ -1822,7 +1822,7 @@ public void snapshotListenerSortsUnicodeStringsInMapAsServer() {
     assertTrue(getSnapshotDocIds.equals(expectedDocIds));
     assertTrue(watchSnapshotDocIds.equals(expectedDocIds));
 
-    checkOnlineAndOfflineResultsMatch(orderedQuery, expectedDocIds.toArray(new String[0]));
+    checkOnlineAndOfflineResultsMatch(colRef, orderedQuery, expectedDocIds.toArray(new String[0]));
   }
 
   @Test
@@ -1879,7 +1879,7 @@ public void snapshotListenerSortsUnicodeStringsInMapKeyAsServer() {
     assertTrue(getSnapshotDocIds.equals(expectedDocIds));
     assertTrue(watchSnapshotDocIds.equals(expectedDocIds));
 
-    checkOnlineAndOfflineResultsMatch(orderedQuery, expectedDocIds.toArray(new String[0]));
+    checkOnlineAndOfflineResultsMatch(colRef, orderedQuery, expectedDocIds.toArray(new String[0]));
   }
 
   @Test
@@ -1937,7 +1937,7 @@ public void snapshotListenerSortsUnicodeStringsInDocumentKeyAsServer() {
     assertTrue(getSnapshotDocIds.equals(expectedDocIds));
     assertTrue(watchSnapshotDocIds.equals(expectedDocIds));
 
-    checkOnlineAndOfflineResultsMatch(orderedQuery, expectedDocIds.toArray(new String[0]));
+    checkOnlineAndOfflineResultsMatch(colRef, orderedQuery, expectedDocIds.toArray(new String[0]));
   }
 
   @Test
@@ -1986,6 +1986,6 @@ public void snapshotListenerSortsInvalidUnicodeStringsAsServer() {
     assertTrue(getSnapshotDocIds.equals(expectedDocIds));
     assertTrue(watchSnapshotDocIds.equals(expectedDocIds));
 
-    checkOnlineAndOfflineResultsMatch(orderedQuery, expectedDocIds.toArray(new String[0]));
+    checkOnlineAndOfflineResultsMatch(colRef, orderedQuery, expectedDocIds.toArray(new String[0]));
   }
 }

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

@@ -1470,30 +1470,38 @@ public void testOrQueries() {
 
     // Two equalities: a==1 || b==1.
     checkOnlineAndOfflineResultsMatch(
-        collection.where(or(equalTo("a", 1), equalTo("b", 1))), "doc1", "doc2", "doc4", "doc5");
+        collection,
+        collection.where(or(equalTo("a", 1), equalTo("b", 1))),
+        "doc1",
+        "doc2",
+        "doc4",
+        "doc5");
 
     // (a==1 && b==0) || (a==3 && b==2)
     checkOnlineAndOfflineResultsMatch(
+        collection,
         collection.where(
             or(and(equalTo("a", 1), equalTo("b", 0)), and(equalTo("a", 3), equalTo("b", 2)))),
         "doc1",
         "doc3");
 
     // a==1 && (b==0 || b==3).
     checkOnlineAndOfflineResultsMatch(
+        collection,
         collection.where(and(equalTo("a", 1), or(equalTo("b", 0), equalTo("b", 3)))),
         "doc1",
         "doc4");
 
     // (a==2 || b==2) && (a==3 || b==3)
     checkOnlineAndOfflineResultsMatch(
+        collection,
         collection.where(
             and(or(equalTo("a", 2), equalTo("b", 2)), or(equalTo("a", 3), equalTo("b", 3)))),
         "doc3");
 
     // Test with limits without orderBy (the __name__ ordering is the tie breaker).
     checkOnlineAndOfflineResultsMatch(
-        collection.where(or(equalTo("a", 2), equalTo("b", 1))).limit(1), "doc2");
+        collection, collection.where(or(equalTo("a", 2), equalTo("b", 1))).limit(1), "doc2");
   }
 
   @Test
@@ -1510,7 +1518,11 @@ public void testOrQueriesWithIn() {
 
     // a==2 || b in [2,3]
     checkOnlineAndOfflineResultsMatch(
-        collection.where(or(equalTo("a", 2), inArray("b", asList(2, 3)))), "doc3", "doc4", "doc6");
+        collection,
+        collection.where(or(equalTo("a", 2), inArray("b", asList(2, 3)))),
+        "doc3",
+        "doc4",
+        "doc6");
   }
 
   @Test
@@ -1527,10 +1539,15 @@ public void testOrQueriesWithArrayMembership() {
 
     // a==2 || b array-contains 7
     checkOnlineAndOfflineResultsMatch(
-        collection.where(or(equalTo("a", 2), arrayContains("b", 7))), "doc3", "doc4", "doc6");
+        collection,
+        collection.where(or(equalTo("a", 2), arrayContains("b", 7))),
+        "doc3",
+        "doc4",
+        "doc6");
 
     // a==2 || b array-contains-any [0, 3]
     checkOnlineAndOfflineResultsMatch(
+        collection,
         collection.where(or(equalTo("a", 2), arrayContainsAny("b", asList(0, 3)))),
         "doc1",
         "doc4",
@@ -1551,12 +1568,12 @@ public void testMultipleInOps() {
 
     // Two IN operations on different fields with disjunction.
     Query query1 = collection.where(or(inArray("a", asList(2, 3)), inArray("b", asList(0, 2))));
-    checkOnlineAndOfflineResultsMatch(query1, "doc1", "doc3", "doc6");
+    checkOnlineAndOfflineResultsMatch(collection, query1, "doc1", "doc3", "doc6");
 
     // Two IN operations on the same field with disjunction.
     // a IN [0,3] || a IN [0,2] should union them (similar to: a IN [0,2,3]).
     Query query2 = collection.where(or(inArray("a", asList(0, 3)), inArray("a", asList(0, 2))));
-    checkOnlineAndOfflineResultsMatch(query2, "doc3", "doc6");
+    checkOnlineAndOfflineResultsMatch(collection, query2, "doc3", "doc6");
   }
 
   @Test
@@ -1573,14 +1590,14 @@ public void testUsingInWithArrayContainsAny() {
 
     Query query1 =
         collection.where(or(inArray("a", asList(2, 3)), arrayContainsAny("b", asList(0, 7))));
-    checkOnlineAndOfflineResultsMatch(query1, "doc1", "doc3", "doc4", "doc6");
+    checkOnlineAndOfflineResultsMatch(collection, query1, "doc1", "doc3", "doc4", "doc6");
 
     Query query2 =
         collection.where(
             or(
                 and(inArray("a", asList(2, 3)), equalTo("c", 10)),
                 arrayContainsAny("b", asList(0, 7))));
-    checkOnlineAndOfflineResultsMatch(query2, "doc1", "doc3", "doc4");
+    checkOnlineAndOfflineResultsMatch(collection, query2, "doc1", "doc3", "doc4");
   }
 
   @Test
@@ -1596,20 +1613,20 @@ public void testUsingInWithArrayContains() {
     CollectionReference collection = testCollectionWithDocs(testDocs);
 
     Query query1 = collection.where(or(inArray("a", asList(2, 3)), arrayContains("b", 3)));
-    checkOnlineAndOfflineResultsMatch(query1, "doc3", "doc4", "doc6");
+    checkOnlineAndOfflineResultsMatch(collection, query1, "doc3", "doc4", "doc6");
 
     Query query2 = collection.where(and(inArray("a", asList(2, 3)), arrayContains("b", 7)));
-    checkOnlineAndOfflineResultsMatch(query2, "doc3");
+    checkOnlineAndOfflineResultsMatch(collection, query2, "doc3");
 
     Query query3 =
         collection.where(
             or(inArray("a", asList(2, 3)), and(arrayContains("b", 3), equalTo("a", 1))));
-    checkOnlineAndOfflineResultsMatch(query3, "doc3", "doc4", "doc6");
+    checkOnlineAndOfflineResultsMatch(collection, query3, "doc3", "doc4", "doc6");
 
     Query query4 =
         collection.where(
             and(inArray("a", asList(2, 3)), or(arrayContains("b", 7), equalTo("a", 1))));
-    checkOnlineAndOfflineResultsMatch(query4, "doc3");
+    checkOnlineAndOfflineResultsMatch(collection, query4, "doc3");
   }
 
   @Test
@@ -1625,9 +1642,9 @@ public void testOrderByEquality() {
     CollectionReference collection = testCollectionWithDocs(testDocs);
 
     Query query1 = collection.where(equalTo("a", 1)).orderBy("a");
-    checkOnlineAndOfflineResultsMatch(query1, "doc1", "doc4", "doc5");
+    checkOnlineAndOfflineResultsMatch(collection, query1, "doc1", "doc4", "doc5");
 
     Query query2 = collection.where(inArray("a", asList(2, 3))).orderBy("a");
-    checkOnlineAndOfflineResultsMatch(query2, "doc6", "doc3");
+    checkOnlineAndOfflineResultsMatch(collection, query2, "doc6", "doc3");
   }
 }

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

@@ -323,7 +323,8 @@ public void vectorFieldOrderOnlineAndOffline() throws Exception {
     Query orderedQuery = randomColl.orderBy("embedding");
 
     // Run query with snapshot listener
-    checkOnlineAndOfflineResultsMatch(orderedQuery, docIds.stream().toArray(String[]::new));
+    checkOnlineAndOfflineResultsMatch(
+        randomColl, orderedQuery, docIds.stream().toArray(String[]::new));
   }
 
   /** Verifies that the SDK filters vector fields the same way for online and offline queries*/
@@ -363,13 +364,15 @@ public void vectorFieldFilterOnlineAndOffline() throws Exception {
             .orderBy("embedding")
             .whereLessThan("embedding", FieldValue.vector(new double[] {1, 2, 100, 4, 4}));
     checkOnlineAndOfflineResultsMatch(
-        orderedQueryLessThan, docIds.subList(2, 11).stream().toArray(String[]::new));
+        randomColl, orderedQueryLessThan, docIds.subList(2, 11).stream().toArray(String[]::new));
 
     Query orderedQueryGreaterThan =
         randomColl
             .orderBy("embedding")
             .whereGreaterThan("embedding", FieldValue.vector(new double[] {1, 2, 100, 4, 4}));
     checkOnlineAndOfflineResultsMatch(
-        orderedQueryGreaterThan, docIds.subList(12, 13).stream().toArray(String[]::new));
+        randomColl,
+        orderedQueryGreaterThan,
+        docIds.subList(12, 13).stream().toArray(String[]::new));
   }
 }

firebase-firestore/src/androidTest/java/com/google/firebase/firestore/testutil/CompositeIndexTestHelper.java

@@ -122,8 +122,10 @@ private Map<String, Map<String, Object>> prepareTestDocuments(
   // actual document IDs created by the test helper.
   @NonNull
   public void assertOnlineAndOfflineResultsMatch(
-      @NonNull Query query, @NonNull String... expectedDocs) {
-    checkOnlineAndOfflineResultsMatch(query, toHashedIds(expectedDocs));
+      @NonNull CollectionReference collection,
+      @NonNull Query query,
+      @NonNull String... expectedDocs) {
+    checkOnlineAndOfflineResultsMatch(collection, query, toHashedIds(expectedDocs));
   }
 
   // Asserts that the IDs in the query snapshot matches the expected Ids. The expected document

firebase-firestore/src/androidTest/java/com/google/firebase/firestore/testutil/IntegrationTestUtil.java

@@ -524,17 +524,41 @@ public static List<Object> nullList() {
    * documents as running the query while offline. If `expectedDocs` is provided, it also checks
    * that both online and offline query result is equal to the expected documents.
    *
+   * This function first performs a "get" for the entire COLLECTION from the server.
+   * It then performs the QUERY from CACHE which, results in `executeFullCollectionScan()`
+   * It then performs the QUERY from SERVER.
+   * It then performs the QUERY from CACHE again, which results in `performQueryUsingRemoteKeys()`.
+   * It then ensure that all the above QUERY results are the same.
+   *
+   * @param collection The collection on which the query is performed.
    * @param query The query to check
    * @param expectedDocs Ordered list of document keys that are expected to match the query
    */
-  public static void checkOnlineAndOfflineResultsMatch(Query query, String... expectedDocs) {
+  public static void checkOnlineAndOfflineResultsMatch(
+      CollectionReference collection, Query query, String... expectedDocs) {
+    // Note: Order matters. The following has to be done in the specific order:
+
+    // 1- Pre-populate the cache with the entire collection.
+    waitFor(collection.get(Source.SERVER));
+
+    // 2- This performs the query against the cache using full collection scan.
+    QuerySnapshot docsFromCacheFullCollectionScan = waitFor(query.get(Source.CACHE));
+
+    // 3- This goes to the server (backend/emulator).
     QuerySnapshot docsFromServer = waitFor(query.get(Source.SERVER));
-    QuerySnapshot docsFromCache = waitFor(query.get(Source.CACHE));
 
-    assertEquals(querySnapshotToIds(docsFromServer), querySnapshotToIds(docsFromCache));
-    List<String> expected = asList(expectedDocs);
-    if (!expected.isEmpty()) {
-      assertEquals(expected, querySnapshotToIds(docsFromCache));
+    // 4- This performs the query against the cache using remote keys.
+    QuerySnapshot docsFromCacheUsingRemoteKeys = waitFor(query.get(Source.CACHE));
+
+    assertEquals(
+        querySnapshotToIds(docsFromServer), querySnapshotToIds(docsFromCacheFullCollectionScan));
+    assertEquals(
+        querySnapshotToIds(docsFromServer), querySnapshotToIds(docsFromCacheUsingRemoteKeys));
+
+    // Expected document IDs.
+    List<String> expectedDocIds = asList(expectedDocs);
+    if (!expectedDocIds.isEmpty()) {
+      assertEquals(expectedDocIds, querySnapshotToIds(docsFromServer));
     }
   }
 }