Firestore: enable use of 'WriteBatch' as a context manager.

firestore/google/cloud/firestore_v1beta1/batch.py

@@ -33,6 +33,8 @@ class WriteBatch(object):
     def __init__(self, client):
         self._client = client
         self._write_pbs = []
+        self.write_results = None
+        self.commit_time = None
 
     def _add_write_pbs(self, write_pbs):
         """Add `Write`` protobufs to this transaction.
@@ -147,4 +149,13 @@ def commit(self):
         )
 
         self._write_pbs = []
-        return list(commit_response.write_results)
+        self.write_results = results = list(commit_response.write_results)
+        self.commit_time = commit_response.commit_time
+        return results
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, exc_type, exc_value, traceback):
+        if exc_type is None:
+            self.commit()

firestore/tests/unit/test_batch.py

@@ -32,6 +32,8 @@ def test_constructor(self):
         batch = self._make_one(mock.sentinel.client)
         self.assertIs(batch._client, mock.sentinel.client)
         self.assertEqual(batch._write_pbs, [])
+        self.assertIsNone(batch.write_results)
+        self.assertIsNone(batch.commit_time)
 
     def test__add_write_pbs(self):
         batch = self._make_one(mock.sentinel.client)
@@ -151,13 +153,16 @@ def test_delete(self):
         self.assertEqual(batch._write_pbs, [new_write_pb])
 
     def test_commit(self):
+        from google.protobuf import timestamp_pb2
         from google.cloud.firestore_v1beta1.proto import firestore_pb2
         from google.cloud.firestore_v1beta1.proto import write_pb2
 
         # Create a minimal fake GAPIC with a dummy result.
         firestore_api = mock.Mock(spec=["commit"])
+        timestamp = timestamp_pb2.Timestamp(seconds=1234567, nanos=123456798)
         commit_response = firestore_pb2.CommitResponse(
-            write_results=[write_pb2.WriteResult(), write_pb2.WriteResult()]
+            write_results=[write_pb2.WriteResult(), write_pb2.WriteResult()],
+            commit_time=timestamp,
         )
         firestore_api.commit.return_value = commit_response
 
@@ -175,6 +180,8 @@ def test_commit(self):
 
         write_results = batch.commit()
         self.assertEqual(write_results, list(commit_response.write_results))
+        self.assertEqual(batch.write_results, write_results)
+        self.assertEqual(batch.commit_time, timestamp)
         # Make sure batch has no more "changes".
         self.assertEqual(batch._write_pbs, [])
 
@@ -186,6 +193,64 @@ def test_commit(self):
             metadata=client._rpc_metadata,
         )
 
+    def test_as_context_mgr_wo_error(self):
+        from google.protobuf import timestamp_pb2
+        from google.cloud.firestore_v1beta1.proto import firestore_pb2
+        from google.cloud.firestore_v1beta1.proto import write_pb2
+
+        firestore_api = mock.Mock(spec=["commit"])
+        timestamp = timestamp_pb2.Timestamp(seconds=1234567, nanos=123456798)
+        commit_response = firestore_pb2.CommitResponse(
+            write_results=[write_pb2.WriteResult(), write_pb2.WriteResult()],
+            commit_time=timestamp,
+        )
+        firestore_api.commit.return_value = commit_response
+        client = _make_client()
+        client._firestore_api_internal = firestore_api
+        batch = self._make_one(client)
+        document1 = client.document("a", "b")
+        document2 = client.document("c", "d", "e", "f")
+
+        with batch as ctx_mgr:
+            self.assertIs(ctx_mgr, batch)
+            ctx_mgr.create(document1, {"ten": 10, "buck": u"ets"})
+            ctx_mgr.delete(document2)
+            write_pbs = batch._write_pbs[::]
+
+        self.assertEqual(batch.write_results, list(commit_response.write_results))
+        self.assertEqual(batch.commit_time, timestamp)
+        # Make sure batch has no more "changes".
+        self.assertEqual(batch._write_pbs, [])
+
+        # Verify the mocks.
+        firestore_api.commit.assert_called_once_with(
+            client._database_string,
+            write_pbs,
+            transaction=None,
+            metadata=client._rpc_metadata,
+        )
+
+    def test_as_context_mgr_w_error(self):
+        firestore_api = mock.Mock(spec=["commit"])
+        client = _make_client()
+        client._firestore_api_internal = firestore_api
+        batch = self._make_one(client)
+        document1 = client.document("a", "b")
+        document2 = client.document("c", "d", "e", "f")
+
+        with self.assertRaises(RuntimeError):
+            with batch as ctx_mgr:
+                ctx_mgr.create(document1, {"ten": 10, "buck": u"ets"})
+                ctx_mgr.delete(document2)
+                raise RuntimeError("testing")
+
+        self.assertIsNone(batch.write_results)
+        self.assertIsNone(batch.commit_time)
+        # batch still has its changes
+        self.assertEqual(len(batch._write_pbs), 2)
+
+        firestore_api.commit.assert_not_called()
+
 
 def _value_pb(**kwargs):
     from google.cloud.firestore_v1beta1.proto.document_pb2 import Value

firestore/tests/unit/test_collection.py

@@ -233,7 +233,9 @@ def test_add_explicit_id(self):
             update_time=mock.sentinel.update_time, spec=["update_time"]
         )
         commit_response = mock.Mock(
-            write_results=[write_result], spec=["write_results"]
+            write_results=[write_result],
+            spec=["write_results", "commit_time"],
+            commit_time=mock.sentinel.commit_time,
         )
         firestore_api.commit.return_value = commit_response
 

firestore/tests/unit/test_document.py

@@ -198,13 +198,19 @@ def _write_pb_for_create(document_path, document_data):
             current_document=common_pb2.Precondition(exists=False),
         )
 
+    @staticmethod
+    def _make_commit_repsonse(write_results=None):
+        from google.cloud.firestore_v1beta1.proto import firestore_pb2
+
+        response = mock.create_autospec(firestore_pb2.CommitResponse)
+        response.write_results = write_results or [mock.sentinel.write_result]
+        response.commit_time = mock.sentinel.commit_time
+        return response
+
     def test_create(self):
         # Create a minimal fake GAPIC with a dummy response.
         firestore_api = mock.Mock(spec=["commit"])
-        commit_response = mock.Mock(
-            write_results=[mock.sentinel.write_result], spec=["write_results"]
-        )
-        firestore_api.commit.return_value = commit_response
+        firestore_api.commit.return_value = self._make_commit_repsonse()
 
         # Attach the fake GAPIC to a real client.
         client = _make_client("dignity")
@@ -235,10 +241,9 @@ def test_create_empty(self):
         snapshot = mock.create_autospec(DocumentSnapshot)
         snapshot.exists = True
         document_reference.get.return_value = snapshot
-        commit_response = mock.Mock(
-            write_results=[document_reference], get=[snapshot], spec=["write_results"]
+        firestore_api.commit.return_value = self._make_commit_repsonse(
+            write_results=[document_reference]
         )
-        firestore_api.commit.return_value = commit_response
 
         # Attach the fake GAPIC to a real client.
         client = _make_client("dignity")
@@ -281,10 +286,7 @@ def _write_pb_for_set(document_path, document_data, merge):
     def _set_helper(self, merge=False, **option_kwargs):
         # Create a minimal fake GAPIC with a dummy response.
         firestore_api = mock.Mock(spec=["commit"])
-        commit_response = mock.Mock(
-            write_results=[mock.sentinel.write_result], spec=["write_results"]
-        )
-        firestore_api.commit.return_value = commit_response
+        firestore_api.commit.return_value = self._make_commit_repsonse()
 
         # Attach the fake GAPIC to a real client.
         client = _make_client("db-dee-bee")
@@ -332,10 +334,7 @@ def _update_helper(self, **option_kwargs):
 
         # Create a minimal fake GAPIC with a dummy response.
         firestore_api = mock.Mock(spec=["commit"])
-        commit_response = mock.Mock(
-            write_results=[mock.sentinel.write_result], spec=["write_results"]
-        )
-        firestore_api.commit.return_value = commit_response
+        firestore_api.commit.return_value = self._make_commit_repsonse()
 
         # Attach the fake GAPIC to a real client.
         client = _make_client("potato-chip")
@@ -389,10 +388,7 @@ def test_update_with_precondition(self):
     def test_empty_update(self):
         # Create a minimal fake GAPIC with a dummy response.
         firestore_api = mock.Mock(spec=["commit"])
-        commit_response = mock.Mock(
-            write_results=[mock.sentinel.write_result], spec=["write_results"]
-        )
-        firestore_api.commit.return_value = commit_response
+        firestore_api.commit.return_value = self._make_commit_repsonse()
 
         # Attach the fake GAPIC to a real client.
         client = _make_client("potato-chip")
@@ -410,10 +406,7 @@ def _delete_helper(self, **option_kwargs):
 
         # Create a minimal fake GAPIC with a dummy response.
         firestore_api = mock.Mock(spec=["commit"])
-        commit_response = mock.Mock(
-            commit_time=mock.sentinel.commit_time, spec=["commit_time"]
-        )
-        firestore_api.commit.return_value = commit_response
+        firestore_api.commit.return_value = self._make_commit_repsonse()
 
         # Attach the fake GAPIC to a real client.
         client = _make_client("donut-base")