Change DocumentOverlayCache.saveOverlays() to require non-null values in the given Map (#3518)

Author: dconeybe

firebase-firestore/src/main/java/com/google/firebase/firestore/local/LocalDocumentsView.java

@@ -202,7 +202,10 @@ private void recalculateAndSaveOverlays(Map<DocumentKey, MutableDocument> docs)
       Map<DocumentKey, Mutation> overlays = new HashMap<>();
       for (DocumentKey key : entry.getValue()) {
         if (!processed.contains(key)) {
-          overlays.put(key, Mutation.calculateOverlayMutation(docs.get(key), masks.get(key)));
+          Mutation mutation = Mutation.calculateOverlayMutation(docs.get(key), masks.get(key));
+          if (mutation != null) {
+            overlays.put(key, mutation);
+          }
           processed.add(key);
         }
       }

firebase-firestore/src/main/java/com/google/firebase/firestore/local/MemoryDocumentOverlayCache.java

@@ -14,6 +14,8 @@
 
 package com.google.firebase.firestore.local;
 
+import static com.google.firebase.firestore.util.Preconditions.checkNotNull;
+
 import androidx.annotation.Nullable;
 import com.google.firebase.firestore.model.DocumentKey;
 import com.google.firebase.firestore.model.ResourcePath;
@@ -50,11 +52,7 @@ public Map<DocumentKey, Overlay> getOverlays(SortedSet<DocumentKey> keys) {
     return result;
   }
 
-  private void saveOverlay(int largestBatchId, @Nullable Mutation mutation) {
-    if (mutation == null) {
-      return;
-    }
-
+  private void saveOverlay(int largestBatchId, Mutation mutation) {
     // Remove the association of the overlay to its batch id.
     Overlay existing = this.overlays.get(mutation.getKey());
     if (existing != null) {
@@ -73,7 +71,8 @@ private void saveOverlay(int largestBatchId, @Nullable Mutation mutation) {
   @Override
   public void saveOverlays(int largestBatchId, Map<DocumentKey, Mutation> overlays) {
     for (Map.Entry<DocumentKey, Mutation> entry : overlays.entrySet()) {
-      saveOverlay(largestBatchId, entry.getValue());
+      Mutation overlay = checkNotNull(entry.getValue(), "null value for key: %s", entry.getKey());
+      saveOverlay(largestBatchId, overlay);
     }
   }
 

firebase-firestore/src/main/java/com/google/firebase/firestore/local/SQLiteDocumentOverlayCache.java

@@ -16,6 +16,7 @@
 
 import static com.google.firebase.firestore.util.Assert.fail;
 import static com.google.firebase.firestore.util.Assert.hardAssert;
+import static com.google.firebase.firestore.util.Preconditions.checkNotNull;
 
 import android.database.Cursor;
 import androidx.annotation.Nullable;
@@ -106,7 +107,7 @@ private void processSingleCollection(
     }
   }
 
-  private void saveOverlay(int largestBatchId, DocumentKey key, @Nullable Mutation mutation) {
+  private void saveOverlay(int largestBatchId, DocumentKey key, Mutation mutation) {
     String group = key.getCollectionGroup();
     String collectionPath = EncodedPath.encode(key.getPath().popLast());
     String documentId = key.getPath().getLastSegment();
@@ -125,9 +126,9 @@ private void saveOverlay(int largestBatchId, DocumentKey key, @Nullable Mutation
   @Override
   public void saveOverlays(int largestBatchId, Map<DocumentKey, Mutation> overlays) {
     for (Map.Entry<DocumentKey, Mutation> entry : overlays.entrySet()) {
-      if (entry.getValue() != null) {
-        saveOverlay(largestBatchId, entry.getKey(), entry.getValue());
-      }
+      DocumentKey key = entry.getKey();
+      Mutation overlay = checkNotNull(entry.getValue(), "null value for key: %s", key);
+      saveOverlay(largestBatchId, key, overlay);
     }
   }
 

firebase-firestore/src/main/java/com/google/firebase/firestore/model/mutation/Mutation.java

@@ -87,6 +87,7 @@ public abstract class Mutation {
    * of the document, and a {@link FieldMask} representing the fields that are mutated by the local
    * mutations.
    */
+  @Nullable
   public static Mutation calculateOverlayMutation(MutableDocument doc, @Nullable FieldMask mask) {
     if ((!doc.hasLocalMutations()) || (mask != null && mask.getMask().isEmpty())) {
       return null;

firebase-firestore/src/main/java/com/google/firebase/firestore/model/mutation/MutationBatch.java

@@ -151,7 +151,9 @@ public Map<DocumentKey, Mutation> applyToLocalDocumentSet(
       // a Set(or Delete) mutation, instead of trying to create a patch mutation as the overlay.
       mutatedFields = documentsWithoutRemoteVersion.contains(key) ? null : mutatedFields;
       Mutation overlay = Mutation.calculateOverlayMutation(document, mutatedFields);
-      overlays.put(key, overlay);
+      if (overlay != null) {
+        overlays.put(key, overlay);
+      }
       if (!document.isValidDocument()) {
         document.convertToNoDocument(SnapshotVersion.NONE);
       }

firebase-firestore/src/main/java/com/google/firebase/firestore/util/Preconditions.java

@@ -14,6 +14,7 @@
 
 package com.google.firebase.firestore.util;
 
+import java.util.Locale;
 import javax.annotation.Nonnull;
 import javax.annotation.Nullable;
 
@@ -149,6 +150,24 @@ public static <T extends Object> T checkNotNull(
     return reference;
   }
 
+  /**
+   * Ensures that an object reference passed as a parameter to the calling method is not null.
+   *
+   * @param reference an object reference
+   * @param errorFormatString the exception message to use if the check fails; will be formatted
+   *     using {@link String#format}.
+   * @param formatArgs the args to specify when formatting the given {@code errorFormatString}.
+   * @return the non-null reference that was validated
+   * @throws NullPointerException if {@code reference} is null
+   */
+  public static <T extends Object> T checkNotNull(
+      @Nonnull T reference, String errorFormatString, @Nullable Object... formatArgs) {
+    if (reference == null) {
+      throw new NullPointerException(String.format(Locale.US, errorFormatString, formatArgs));
+    }
+    return reference;
+  }
+
   /**
    * Ensures the truth of an expression involving the state of the calling instance, but not
    * involving any parameters to the calling method.

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

@@ -23,6 +23,7 @@
 import static com.google.firebase.firestore.testutil.TestUtil.setMutation;
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertNull;
+import static org.junit.Assert.assertThrows;
 import static org.junit.Assert.assertTrue;
 
 import com.google.firebase.firestore.auth.User;
@@ -41,6 +42,7 @@
 import org.junit.After;
 import org.junit.Before;
 import org.junit.Test;
+import org.junit.function.ThrowingRunnable;
 
 /**
  * These are tests for any implementation of the DocumentOverlayCache interface.
@@ -252,6 +254,24 @@ public void testUpdateDocumentOverlay() {
     assertNull(cache.getOverlay(DocumentKey.fromPathString("coll/doc")));
   }
 
+  @Test
+  public void testSaveOverlaysThrowsNullPointerExceptionOnNullMapValue() {
+    Map<DocumentKey, Mutation> data = new HashMap<>();
+    data.put(key("coll/doc"), null);
+
+    NullPointerException e =
+        assertThrows(
+            NullPointerException.class,
+            new ThrowingRunnable() {
+              @Override
+              public void run() {
+                cache.saveOverlays(1, data);
+              }
+            });
+
+    assertThat(e.getMessage()).contains("coll/doc");
+  }
+
   void verifyOverlayContains(Map<DocumentKey, Overlay> overlays, String... keys) {
     Set<DocumentKey> expected = Arrays.stream(keys).map(TestUtil::key).collect(Collectors.toSet());
     assertThat(overlays.keySet()).containsExactlyElementsIn(expected);