Merge Option True

Author: chemelnucfin

firestore/google/cloud/firestore_v1beta1/client.py

@@ -288,7 +288,7 @@ def write_option(**kwargs):
         elif name == 'exists':
             return ExistsOption(value)
         elif name == 'merge':
-            return MergeOption()
+            return MergeOption(value)
         else:
             extra = '{!r} was provided'.format(name)
             raise TypeError(_BAD_OPTION_ERR, extra)
@@ -429,7 +429,15 @@ class MergeOption(WriteOption):
 
     This will typically be created by
     :meth:`~.firestore_v1beta1.client.Client.write_option`.
+
+    Args:
+        merge (bool): 
+            The only valid option is True. Any other argument will make this
+            option ignored.
     """
+    def __init__(self, merge):
+        self._merge = merge
+        
     def modify_write(
             self, write_pb, field_paths=None, path=None, **unused_kwargs):
         """Modify a ``Write`` protobuf based on the state of this write option.
@@ -445,9 +453,10 @@ def modify_write(
             unused_kwargs (Dict[str, Any]): Keyword arguments accepted by
                 other subclasses that are unused here.
         """
-        field_paths = sorted(field_paths)  # for testing purposes
-        mask = common_pb2.DocumentMask(field_paths=field_paths)
-        write_pb.update_mask.CopyFrom(mask)
+        if self._merge is True:
+            field_paths = sorted(field_paths)  # for testing purposes
+            mask = common_pb2.DocumentMask(field_paths=field_paths)
+            write_pb.update_mask.CopyFrom(mask)
 
 
 class ExistsOption(WriteOption):

firestore/tests/unit/test_batch.py

@@ -103,7 +103,7 @@ def test_set_merge(self):
         field = 'zapzap'
         value = u'meadows and flowers'
         document_data = {field: value}
-        option = MergeOption()
+        option = MergeOption(True)
         ret_val = batch.set(reference, document_data, option)
         self.assertIsNone(ret_val)
         new_write_pb = write_pb2.Write(

firestore/tests/unit/test_client.py

@@ -489,21 +489,30 @@ def _get_target_class():
 
     def _make_one(self, *args, **kwargs):
         klass = self._get_target_class()
-        return klass()
+        return klass(*args)
 
-    def test_modify_write(self):
+    def test_modify_write_true(self):
         from google.cloud.firestore_v1beta1.proto import common_pb2
         from google.cloud.firestore_v1beta1.proto import write_pb2
 
-        for merge in (True, False):
-            option = self._make_one(merge)
-            write_pb = write_pb2.Write()
-            field_paths = ['a', 'b', 'c']
-            ret_val = option.modify_write(write_pb, field_paths=field_paths)
-            mask = common_pb2.DocumentMask(field_paths=sorted(field_paths))
-            self.assertIsNone(ret_val)
-            self.assertEqual(write_pb.update_mask, mask)
+        option = self._make_one(True)
+        write_pb = write_pb2.Write()
+        field_paths = ['a', 'b', 'c']
+        ret_val = option.modify_write(write_pb, field_paths=field_paths)
+        mask = common_pb2.DocumentMask(field_paths=sorted(field_paths))
+        self.assertIsNone(ret_val)
+        self.assertEqual(write_pb.update_mask, mask)
 
+    def test_modify_write_false(self):
+        from google.cloud.firestore_v1beta1.proto import common_pb2
+        from google.cloud.firestore_v1beta1.proto import write_pb2
+
+        option = self._make_one(False)
+        write_pb = write_pb2.Write()
+        field_paths = ['a', 'b', 'c']
+        ret_val = option.modify_write(write_pb, field_paths=field_paths)
+        mask = common_pb2.DocumentMask(field_paths=sorted(field_paths))
+        self.assertIsNone(ret_val)
 
 class Test__reference_info(unittest.TestCase):
 

firestore/tests/unit/test_document.py

@@ -297,7 +297,7 @@ def test_set_exists(self):
         self._set_helper(exists=True)
 
     def test_set_merge(self):
-        self._set_helper(merge='abc')
+        self._set_helper(merge=True)
 
     @staticmethod
     def _write_pb_for_update(document_path, update_values, field_paths):