From 7b8d9b0e6beff63ff5dcefd10daf77b036a57ead Mon Sep 17 00:00:00 2001
From: Tres Seaver <tseaver@palladion.com>
Date: Wed, 7 Oct 2020 13:13:56 -0400
Subject: [PATCH] feat: add 'update_transforms' to '_helpers.pbs_for_update'

Toward #217.
---
 google/cloud/firestore_v1/_helpers.py | 20 ++++++--------------
 tests/unit/v1/test__helpers.py        | 22 +++++++++-------------
 2 files changed, 15 insertions(+), 27 deletions(-)

diff --git a/google/cloud/firestore_v1/_helpers.py b/google/cloud/firestore_v1/_helpers.py
index d68ec7ba0..e98ec8547 100644
--- 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
index 56dbc3287..488f45fd5 100644
--- 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