From c643d914075c1bfc2549a56ec419aff90af4d8e7 Mon Sep 17 00:00:00 2001 From: Chris McDonough Date: Fri, 11 Jan 2019 16:04:43 -0500 Subject: [PATCH] fix: respect transform values passed into collection.add (#7072) closes #6826 respect transform values passed into collection.add --- google/cloud/firestore_v1beta1/collection.py | 8 ++--- tests/system.py | 2 -- tests/unit/test_collection.py | 32 +++++++++++++++----- 3 files changed, 28 insertions(+), 14 deletions(-) diff --git a/google/cloud/firestore_v1beta1/collection.py b/google/cloud/firestore_v1beta1/collection.py index 9d616fe4d..9c0f98ac7 100644 --- a/google/cloud/firestore_v1beta1/collection.py +++ b/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) diff --git a/tests/system.py b/tests/system.py index 226b1bd9b..670b3dcdf 100644 --- a/tests/system.py +++ b/tests/system.py @@ -409,7 +409,6 @@ def test_collection_add(client, cleanup): cleanup(document_ref1) snapshot1 = document_ref1.get() assert snapshot1.to_dict() == data1 - assert snapshot1.create_time == update_time1 assert snapshot1.update_time == update_time1 assert RANDOM_ID_REGEX.match(document_ref1.id) @@ -429,7 +428,6 @@ def test_collection_add(client, cleanup): cleanup(document_ref3) snapshot3 = document_ref3.get() assert snapshot3.to_dict() == data3 - assert snapshot3.create_time == update_time3 assert snapshot3.update_time == update_time3 assert RANDOM_ID_REGEX.match(document_ref3.id) diff --git a/tests/unit/test_collection.py b/tests/unit/test_collection.py index 6d555526e..09fa1ffe2 100644 --- a/tests/unit/test_collection.py +++ b/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):