feat: add 'update_transforms' to '_helpers.pbs_for_update'

Toward #217.
googleapis · Oct 7, 2020 · 7b8d9b0 · 7b8d9b0
1 parent 0c2f950
commit 7b8d9b0
Show file tree

Hide file tree

Showing 2 changed files with 15 additions and 27 deletions.
diff --git a/google/cloud/firestore_v1/_helpers.py b/google/cloud/firestore_v1/_helpers.py
@@ -866,22 +866,14 @@ def pbs_for_update(document_path, field_updates, option) -> List[types.write.Wri
     if option is None:  # Default is to use ``exists=True``.
         option = ExistsOption(exists=True)
 
-    write_pbs = []
-
-    if extractor.field_paths or extractor.deleted_fields:
-        update_pb = extractor.get_update_pb(document_path)
-        option.modify_write(update_pb)
-        write_pbs.append(update_pb)
+    update_pb = extractor.get_update_pb(document_path)
+    option.modify_write(update_pb)
 
     if extractor.has_transforms:
-        transform_pb = extractor.get_transform_pb(document_path)
-        if not write_pbs:
-            # NOTE: set the write option on the ``transform_pb`` only if there
-            #       is no ``update_pb``
-            option.modify_write(transform_pb)
-        write_pbs.append(transform_pb)
-
-    return write_pbs
+        field_transform_pbs = extractor.get_field_transform_pbs(document_path)
+        update_pb.update_transforms.extend(field_transform_pbs)
+
+    return [update_pb]
 
 
 def pb_for_delete(document_path, option) -> types.write.Write:

diff --git a/tests/unit/v1/test__helpers.py b/tests/unit/v1/test__helpers.py
@@ -2159,23 +2159,19 @@ def _helper(self, option=None, do_transform=False, **write_kwargs):
         if isinstance(option, _helpers.ExistsOption):
             precondition = common.Precondition(exists=False)
             expected_update_pb._pb.current_document.CopyFrom(precondition._pb)
-        expected_pbs = [expected_update_pb]
+
         if do_transform:
             transform_paths = FieldPath.from_string(field_path2)
             server_val = DocumentTransform.FieldTransform.ServerValue
-            expected_transform_pb = write.Write(
-                transform=write.DocumentTransform(
-                    document=document_path,
-                    field_transforms=[
-                        write.DocumentTransform.FieldTransform(
-                            field_path=transform_paths.to_api_repr(),
-                            set_to_server_value=server_val.REQUEST_TIME,
-                        )
-                    ],
+            field_transform_pbs = [
+                write.DocumentTransform.FieldTransform(
+                    field_path=transform_paths.to_api_repr(),
+                    set_to_server_value=server_val.REQUEST_TIME,
                 )
-            )
-            expected_pbs.append(expected_transform_pb)
-        self.assertEqual(write_pbs, expected_pbs)
+            ]
+            expected_update_pb.update_transforms.extend(field_transform_pbs)
+
+        self.assertEqual(write_pbs, [expected_update_pb])
 
     def test_without_option(self):
         from google.cloud.firestore_v1.types import common