From 5e5808d6372e03b7f27ca17d0248be0afe804ba1 Mon Sep 17 00:00:00 2001
From: cherylEnkidu <cheryllin@google.com>
Date: Sat, 24 Jun 2023 14:56:24 -0400
Subject: [PATCH] add mutation to local test

---
 .../local/AutoIndexingExperiment.java         | 30 +++++++++++++++++++
 1 file changed, 30 insertions(+)

diff --git a/firebase-firestore/src/test/java/com/google/firebase/firestore/local/AutoIndexingExperiment.java b/firebase-firestore/src/test/java/com/google/firebase/firestore/local/AutoIndexingExperiment.java
index db22a1b66f0c..56b4bdedb7f7 100644
--- a/firebase-firestore/src/test/java/com/google/firebase/firestore/local/AutoIndexingExperiment.java
+++ b/firebase-firestore/src/test/java/com/google/firebase/firestore/local/AutoIndexingExperiment.java
@@ -4,10 +4,12 @@
 import static com.google.firebase.firestore.testutil.TestUtil.docMap;
 import static com.google.firebase.firestore.testutil.TestUtil.filter;
 import static com.google.firebase.firestore.testutil.TestUtil.map;
+import static com.google.firebase.firestore.testutil.TestUtil.patchMutation;
 import static com.google.firebase.firestore.testutil.TestUtil.query;
 import static org.junit.Assert.assertEquals;
 
 import com.google.android.gms.common.internal.Preconditions;
+import com.google.firebase.Timestamp;
 import com.google.firebase.database.collection.ImmutableSortedMap;
 import com.google.firebase.database.collection.ImmutableSortedSet;
 import com.google.firebase.firestore.auth.User;
@@ -18,6 +20,8 @@
 import com.google.firebase.firestore.model.DocumentSet;
 import com.google.firebase.firestore.model.FieldIndex;
 import com.google.firebase.firestore.model.MutableDocument;
+import com.google.firebase.firestore.model.mutation.Mutation;
+import com.google.firebase.firestore.model.mutation.MutationBatch;
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.Collections;
@@ -99,6 +103,19 @@ protected void addDocument(MutableDocument... docs) {
         });
   }
 
+  protected void addMutation(Mutation mutation) {
+    persistence.runTransaction(
+        "addMutation",
+        () -> {
+          MutationBatch batch =
+              mutationQueue.addMutationBatch(
+                  Timestamp.now(), Collections.emptyList(), Collections.singletonList(mutation));
+          Map<DocumentKey, Mutation> overlayMap =
+              Collections.singletonMap(mutation.getKey(), mutation);
+          documentOverlayCache.saveOverlays(batch.getBatchId(), overlayMap);
+        });
+  }
+
   protected <T> T expectOptimizedCollectionScan(Callable<T> c) throws Exception {
     try {
       expectFullCollectionScan = false;
@@ -169,6 +186,18 @@ private void createTestingCollection(
     }
   }
 
+  private void createMutationForCollection(String basePath, int totalSetCount) {
+    ArrayList<Integer> indexes = new ArrayList<>();
+    for (int index = 0; index < totalSetCount * 10; index++) {
+      indexes.add(index);
+    }
+    Collections.shuffle(indexes);
+
+    for (int i = 0; i < totalSetCount; i++) {
+      addMutation(patchMutation(basePath + "/" + indexes.get(i), map("a", 5)));
+    }
+  }
+
   @Test
   public void testCombinesIndexedWithNonIndexedResults() throws Exception {
     // Every set contains 10 documents
@@ -192,6 +221,7 @@ public void testCombinesIndexedWithNonIndexedResults() throws Exception {
           Query query = query(basePath).filter(filter("match", "==", true));
           indexManager.createTargetIndices(query.toTarget());
           createTestingCollection(basePath, totalSetCount, portion, numOfFields);
+          createMutationForCollection(basePath, totalSetCount);
 
           QueryContext counterWithoutIndex = new QueryContext();
           long beforeAutoStart = System.nanoTime();