fix(firestore): simplify 'Collection.add', avoid spurious API call (#9634)

tseaver · web-flow · commit ac0c96b063a9 · 2019-11-07T15:03:20.000-05:00
Closes #9629
diff --git a/packages/google-cloud-firestore/google/cloud/firestore_v1/collection.py b/packages/google-cloud-firestore/google/cloud/firestore_v1/collection.py
@@ -20,7 +20,6 @@
 
 from google.cloud.firestore_v1 import _helpers
 from google.cloud.firestore_v1 import query as query_mod
-from google.cloud.firestore_v1.proto import document_pb2
 from google.cloud.firestore_v1.watch import Watch
 from google.cloud.firestore_v1 import document
 
@@ -157,27 +156,11 @@ def add(self, document_data, document_id=None):
                 and the document already exists.
         """
         if document_id is None:
-            parent_path, expected_prefix = self._parent_info()
-
-            document_pb = document_pb2.Document()
-
-            created_document_pb = self._client._firestore_api.create_document(
-                parent_path,
-                collection_id=self.id,
-                document_id=None,
-                document=document_pb,
-                mask=None,
-                metadata=self._client._rpc_metadata,
-            )
+            document_id = _auto_id()
 
-            new_document_id = _helpers.get_doc_id(created_document_pb, expected_prefix)
-            document_ref = self.document(new_document_id)
-            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)
-            return write_result.update_time, document_ref
+        document_ref = self.document(document_id)
+        write_result = document_ref.create(document_data)
+        return write_result.update_time, document_ref
 
     def list_documents(self, page_size=None):
         """List all subdocuments of the current collection.
diff --git a/packages/google-cloud-firestore/tests/unit/v1/test_collection.py b/packages/google-cloud-firestore/tests/unit/v1/test_collection.py
@@ -12,7 +12,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import datetime
 import types
 import unittest
 
@@ -193,7 +192,7 @@ def test_add_auto_assigned(self):
         from google.cloud.firestore_v1.proto import document_pb2
         from google.cloud.firestore_v1.document import DocumentReference
         from google.cloud.firestore_v1 import SERVER_TIMESTAMP
-        from google.cloud.firestore_v1._helpers import pbs_for_set_no_merge
+        from google.cloud.firestore_v1._helpers import pbs_for_create
 
         # Create a minimal fake GAPIC add attach it to a real client.
         firestore_api = mock.Mock(spec=["create_document", "commit"])
@@ -214,42 +213,32 @@ def test_add_auto_assigned(self):
         # Actually make a collection.
         collection = self._make_one("grand-parent", "parent", "child", client=client)
 
-        # Add a dummy response for the fake GAPIC.
-        parent_path = collection.parent._document_path
-        auto_assigned_id = "cheezburger"
-        name = "{}/{}/{}".format(parent_path, collection.id, auto_assigned_id)
-        create_doc_response = document_pb2.Document(name=name)
-        create_doc_response.update_time.FromDatetime(datetime.datetime.utcnow())
-        firestore_api.create_document.return_value = create_doc_response
-
         # 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)
+
+        patch = mock.patch("google.cloud.firestore_v1.collection._auto_id")
+        random_doc_id = "DEADBEEF"
+        with patch as patched:
+            patched.return_value = random_doc_id
+            update_time, document_ref = collection.add(document_data)
 
         # Verify the response and the mocks.
         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,)
+        expected_path = collection._path + (random_doc_id,)
         self.assertEqual(document_ref._path, expected_path)
 
-        expected_document_pb = document_pb2.Document()
-        firestore_api.create_document.assert_called_once_with(
-            parent_path,
-            collection_id=collection.id,
-            document_id=None,
-            document=expected_document_pb,
-            mask=None,
-            metadata=client._rpc_metadata,
-        )
-        write_pbs = pbs_for_set_no_merge(document_ref._document_path, document_data)
+        write_pbs = pbs_for_create(document_ref._document_path, document_data)
         firestore_api.commit.assert_called_once_with(
             client._database_string,
             write_pbs,
             transaction=None,
             metadata=client._rpc_metadata,
         )
+        # Since we generate the ID locally, we don't call 'create_document'.
+        firestore_api.create_document.assert_not_called()
 
     @staticmethod
     def _write_pb_for_create(document_path, document_data):
