fix: respect transform values passed into collection.add

firestore/google/cloud/firestore_v1beta1/collection.py

@@ -159,9 +159,8 @@ def add(self, document_data, document_id=None):
         """
         if document_id is None:
             parent_path, expected_prefix = self._parent_info()
-            document_pb = document_pb2.Document(
-                fields=_helpers.encode_dict(document_data)
-            )
+
+            document_pb = document_pb2.Document()
 
             created_document_pb = self._client._firestore_api.create_document(
                 parent_path,
@@ -174,7 +173,8 @@ def add(self, document_data, document_id=None):
 
             new_document_id = _helpers.get_doc_id(created_document_pb, expected_prefix)
             document_ref = self.document(new_document_id)
-            return created_document_pb.update_time, document_ref
+            set_result = document_ref.set(document_data)
+            return set_result.update_time, document_ref
         else:
             document_ref = self.document(document_id)
             write_result = document_ref.create(document_data)

firestore/tests/system.py

@@ -409,7 +409,13 @@ def test_collection_add(client, cleanup):
     cleanup(document_ref1)
     snapshot1 = document_ref1.get()
     assert snapshot1.to_dict() == data1
-    assert snapshot1.create_time == update_time1
+    # NB: We used to assert here that the snapshot create_time was the document
+    # update time as returned by add(), but now under the hood, to allow for
+    # transform values passed into add, the document is created, then a set()
+    # is issued against the document later; the value returned by add is the
+    # result of the latter.  We don't have access to the create time,
+    # so we don't have a value to make this assertion with.
+    # assert snapshot1.create_time == update_time1
     assert snapshot1.update_time == update_time1
     assert RANDOM_ID_REGEX.match(document_ref1.id)
 
@@ -429,7 +435,8 @@ def test_collection_add(client, cleanup):
     cleanup(document_ref3)
     snapshot3 = document_ref3.get()
     assert snapshot3.to_dict() == data3
-    assert snapshot3.create_time == update_time3
+    # see NB above
+    #assert snapshot3.create_time == update_time3
     assert snapshot3.update_time == update_time3
     assert RANDOM_ID_REGEX.match(document_ref3.id)
 

firestore/tests/unit/test_collection.py

@@ -191,11 +191,21 @@ def test__parent_info_nested(self):
 
     def test_add_auto_assigned(self):
         from google.cloud.firestore_v1beta1.proto import document_pb2
-        from google.cloud.firestore_v1beta1 import _helpers
         from google.cloud.firestore_v1beta1.document import DocumentReference
+        from google.cloud.firestore_v1beta1 import SERVER_TIMESTAMP
+        from google.cloud.firestore_v1beta1._helpers import pbs_for_set_no_merge
 
         # Create a minimal fake GAPIC add attach it to a real client.
-        firestore_api = mock.Mock(spec=["create_document"])
+        firestore_api = mock.Mock(spec=["create_document", "commit"])
+        write_result = mock.Mock(
+            update_time=mock.sentinel.update_time, spec=["update_time"]
+        )
+        commit_response = mock.Mock(
+            write_results=[write_result],
+            spec=["write_results", "commit_time"],
+            commit_time=mock.sentinel.commit_time,
+        )
+        firestore_api.commit.return_value = commit_response
         create_doc_response = document_pb2.Document()
         firestore_api.create_document.return_value = create_doc_response
         client = _make_client()
@@ -212,20 +222,19 @@ def test_add_auto_assigned(self):
         create_doc_response.update_time.FromDatetime(datetime.datetime.utcnow())
         firestore_api.create_document.return_value = create_doc_response
 
-        # Actually call add() on our collection.
-        document_data = {"been": "here"}
+        # Actually call add() on our collection; include a transform to make
+        # sure transforms during adds work.
+        document_data = {"been": "here", "now": SERVER_TIMESTAMP}
         update_time, document_ref = collection.add(document_data)
 
         # Verify the response and the mocks.
-        self.assertIs(update_time, create_doc_response.update_time)
+        self.assertIs(update_time, mock.sentinel.update_time)
         self.assertIsInstance(document_ref, DocumentReference)
         self.assertIs(document_ref._client, client)
         expected_path = collection._path + (auto_assigned_id,)
         self.assertEqual(document_ref._path, expected_path)
 
-        expected_document_pb = document_pb2.Document(
-            fields=_helpers.encode_dict(document_data)
-        )
+        expected_document_pb = document_pb2.Document()
         firestore_api.create_document.assert_called_once_with(
             parent_path,
             collection_id=collection.id,
@@ -234,6 +243,13 @@ def test_add_auto_assigned(self):
             mask=None,
             metadata=client._rpc_metadata,
         )
+        write_pbs = pbs_for_set_no_merge(document_ref._document_path, document_data)
+        firestore_api.commit.assert_called_once_with(
+            client._database_string,
+            write_pbs,
+            transaction=None,
+            metadata=client._rpc_metadata,
+        )
 
     @staticmethod
     def _write_pb_for_create(document_path, document_data):