add extra test helper, move terraform into firestore folder

Author: milaGGL

.github/workflows/ci_tests.yml

@@ -127,11 +127,17 @@ jobs:
         uses: hashicorp/setup-terraform@v2
       - name: Terraform Init
         if: contains(matrix.module, ':firebase-firestore')
-        run: terraform init
+        run: |
+          pwd
+          cd firebase-firestore
+          terraform init
         continue-on-error: true
       - name: Terraform Apply
         if: github.event_name == 'pull_request' && contains(matrix.module, ':firebase-firestore')
-        run: terraform apply -var-file=google-services.json -auto-approve
+        run: |
+          pwd
+          cd firebase-firestore
+          terraform apply -var-file=../google-services.json -auto-approve
         continue-on-error: true
       - name: ${{ matrix.module }} Integ Tests
         env:

firestore_index_config.tf renamed to firebase-firestore/firestore_index_config.tf

@@ -59,37 +59,25 @@ public void testOrQueriesWithCompositeIndexes() {
             "doc5", map("a", 1, "b", 1));
     CollectionReference collection = testHelper.withTestDocs(testDocs);
 
+    Query query = collection.where(or(greaterThan("a", 2), equalTo("b", 1)));
     // with one inequality: a>2 || b==1.
-    testHelper.assertOnlineAndOfflineResultsMatch(
-        testHelper.query(collection.where(or(greaterThan("a", 2), equalTo("b", 1)))),
-        "doc5",
-        "doc2",
-        "doc3");
+    testHelper.assertOnlineAndOfflineResultsMatch(testHelper.query(query), "doc5", "doc2", "doc3");
 
     // Test with limits (implicit order by ASC): (a==1) || (b > 0) LIMIT 2
-    testHelper.assertOnlineAndOfflineResultsMatch(
-        testHelper.query(collection.where(or(equalTo("a", 1), greaterThan("b", 0))).limit(2)),
-        "doc1",
-        "doc2");
+    query = collection.where(or(equalTo("a", 1), greaterThan("b", 0))).limit(2);
+    testHelper.assertOnlineAndOfflineResultsMatch(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.
-    testHelper.assertOnlineAndOfflineResultsMatch(
-        testHelper.query(
-            collection.where(or(equalTo("a", 1), greaterThan("b", 0))).limitToLast(2).orderBy("b")),
-        "doc3",
-        "doc4");
+    query = collection.where(or(equalTo("a", 1), greaterThan("b", 0))).limitToLast(2).orderBy("b");
+    testHelper.assertOnlineAndOfflineResultsMatch(testHelper.query(query), "doc3", "doc4");
 
     // Test with limits (explicit order by ASC): (a==2) || (b == 1) ORDER BY a LIMIT 1
-    testHelper.assertOnlineAndOfflineResultsMatch(
-        testHelper.query(
-            collection.where(or(equalTo("a", 2), equalTo("b", 1))).limit(1).orderBy("a")),
-        "doc5");
+    query = collection.where(or(equalTo("a", 2), equalTo("b", 1))).limit(1).orderBy("a");
+    testHelper.assertOnlineAndOfflineResultsMatch(testHelper.query(query), "doc5");
 
     // Test with limits (explicit order by DESC): (a==2) || (b == 1) ORDER BY a LIMIT_TO_LAST 1
-    testHelper.assertOnlineAndOfflineResultsMatch(
-        testHelper.query(
-            collection.where(or(equalTo("a", 2), equalTo("b", 1))).limitToLast(1).orderBy("a")),
-        "doc2");
+    query = collection.where(or(equalTo("a", 2), equalTo("b", 1))).limitToLast(1).orderBy("a");
+    testHelper.assertOnlineAndOfflineResultsMatch(testHelper.query(query), "doc2");
   }
 }

main.tf renamed to firebase-firestore/main.tf

@@ -16,6 +16,7 @@
 
 import static com.google.firebase.firestore.testutil.IntegrationTestUtil.checkOnlineAndOfflineResultsMatch;
 import static com.google.firebase.firestore.testutil.IntegrationTestUtil.testFirestore;
+import static com.google.firebase.firestore.testutil.IntegrationTestUtil.waitFor;
 import static com.google.firebase.firestore.testutil.IntegrationTestUtil.writeAllDocs;
 import static com.google.firebase.firestore.util.Util.autoId;
 
@@ -24,21 +25,31 @@
 import com.google.firebase.Timestamp;
 import com.google.firebase.firestore.CollectionReference;
 import com.google.firebase.firestore.DocumentReference;
+import com.google.firebase.firestore.DocumentSnapshot;
 import com.google.firebase.firestore.Query;
+import com.google.firebase.firestore.QuerySnapshot;
+import java.util.ArrayList;
 import java.util.HashMap;
+import java.util.List;
 import java.util.Map;
 
 /**
  * This helper class is designed to facilitate integration testing of Firestore queries that require
  * composite indexes within a controlled testing environment.
  *
- * <p>Key Features: - Runs tests against the dedicated test collection with predefined composite
- * indexes. - Automatically associates a test ID with documents for data isolation. - Utilizes TTL
- * policy for automatic test data cleanup. - Constructs Firestore queries with test ID filters.
+ * <p>Key Features:
+ *
+ * <ul>
+ *   <li>Runs tests against the dedicated test collection with predefined composite indexes.
+ *   <li>Automatically associates a test ID with documents for data isolation.
+ *   <li>Utilizes TTL policy for automatic test data cleanup.
+ *   <li>Constructs Firestore queries with test ID filters.
+ * </ul>
  */
 public class CompositeIndexTestHelper {
   private final String testId;
   private static final String TEST_ID_FIELD = "testId";
+  private static final String TTL_FIELD = "expireAt";
   private static final String COMPOSITE_INDEX_TEST_COLLECTION = "composite-index-test-collection";
 
   // Creates a new instance of the CompositeIndexTestHelper class, with a unique test
@@ -56,12 +67,6 @@ public CollectionReference withTestDocs(@NonNull Map<String, Map<String, Object>
     return reader;
   }
 
-  // Adds a filter on test id for a query.
-  @NonNull
-  public Query query(@NonNull Query query_) {
-    return query_.whereEqualTo(TEST_ID_FIELD, testId);
-  }
-
   // Hash the document key with testId.
   private String toHashedId(String docId) {
     return docId + '-' + testId;
@@ -80,12 +85,19 @@ private Map<String, Object> addTestSpecificFieldsToDoc(Map<String, Object> doc)
     Map<String, Object> updatedDoc = new HashMap<>(doc);
     updatedDoc.put(TEST_ID_FIELD, testId);
     updatedDoc.put(
-        "expireAt",
+        TTL_FIELD,
         new Timestamp( // Expire test data after 24 hours
             Timestamp.now().getSeconds() + 24 * 60 * 60, Timestamp.now().getNanoseconds()));
     return updatedDoc;
   }
 
+  // Remove test-specific fields from a document.
+  private Map<String, Object> removeTestSpecificFieldsFromDoc(Map<String, Object> doc) {
+    doc.remove(TTL_FIELD);
+    doc.remove(TEST_ID_FIELD);
+    return doc;
+  }
+
   // Helper method to hash document keys and add test-specific fields for the provided documents.
   private Map<String, Map<String, Object>> prepareTestDocuments(
       Map<String, Map<String, Object>> docs) {
@@ -106,6 +118,22 @@ public void assertOnlineAndOfflineResultsMatch(
     checkOnlineAndOfflineResultsMatch(query, toHashedIds(expectedDocs));
   }
 
+  // Adds a filter on test id for a query.
+  @NonNull
+  public Query query(@NonNull Query query_) {
+    return query_.whereEqualTo(TEST_ID_FIELD, testId);
+  }
+
+  // Get document reference from a document key.
+  @NonNull
+  public DocumentReference getDocRef(
+      @NonNull CollectionReference collection, @NonNull String docId) {
+    if (!docId.contains("test-id-")) {
+      docId = toHashedId(docId);
+    }
+    return collection.document(docId);
+  }
+
   // Adds a document to a Firestore collection with test-specific fields.
   @NonNull
   public Task<DocumentReference> addDoc(
@@ -118,4 +146,33 @@ public Task<DocumentReference> addDoc(
   public Task<Void> setDoc(@NonNull DocumentReference document, @NonNull Map<String, Object> data) {
     return document.set(addTestSpecificFieldsToDoc(data));
   }
+
+  @NonNull
+  public Task<Void> updateDoc(
+      @NonNull DocumentReference document, @NonNull Map<String, Object> data) {
+    return document.update(data);
+  }
+
+  @NonNull
+  public Task<Void> deleteDoc(@NonNull DocumentReference document) {
+    return document.delete();
+  }
+
+  // Retrieves a single document from Firestore with test-specific fields removed.
+  @NonNull
+  public Map<String, Object> getDoc(@NonNull DocumentReference document) {
+    DocumentSnapshot docSnapshot = waitFor(document.get());
+    return removeTestSpecificFieldsFromDoc(docSnapshot.getData());
+  }
+
+  // Retrieves multiple documents from Firestore with test-specific fields removed.
+  @NonNull
+  public List<Map<String, Object>> getDocs(@NonNull Query query_) {
+    QuerySnapshot querySnapshot = waitFor(query(query_).get());
+    List<Map<String, Object>> res = new ArrayList<>();
+    for (DocumentSnapshot doc : querySnapshot) {
+      res.add(removeTestSpecificFieldsFromDoc(doc.getData()));
+    }
+    return res;
+  }
 }