From 559f2ebff6eb224823fba7d12dcfc2b902b1f71c Mon Sep 17 00:00:00 2001
From: Tres Seaver <tseaver@palladion.com>
Date: Tue, 13 Oct 2020 16:52:17 -0400
Subject: [PATCH] feat: add retry/timeout to 'document.DocumentReference.set'

Toward #221
---
 google/cloud/firestore_v1/document.py | 14 ++++++++++++--
 tests/unit/v1/test_document.py        | 21 +++++++++++++++++++--
 2 files changed, 31 insertions(+), 4 deletions(-)

diff --git a/google/cloud/firestore_v1/document.py b/google/cloud/firestore_v1/document.py
index de2cc4ec1..f56794db8 100644
--- a/google/cloud/firestore_v1/document.py
+++ b/google/cloud/firestore_v1/document.py
@@ -96,7 +96,13 @@ def create(
         write_results = batch.commit(**kwargs)
         return _first_write_result(write_results)
 
-    def set(self, document_data: dict, merge: bool = False) -> Any:
+    def set(
+        self,
+        document_data: dict,
+        merge: bool = False,
+        retry: retries.Retry = None,
+        timeout: float = None,
+    ) -> Any:
         """Replace the current document in the Firestore database.
 
         A write ``option`` can be specified to indicate preconditions of
@@ -116,6 +122,9 @@ def set(self, document_data: dict, merge: bool = False) -> Any:
             merge (Optional[bool] or Optional[List<apispec>]):
                 If True, apply merging instead of overwriting the state
                 of the document.
+            retry (google.api_core.retry.Retry): Designation of what errors, if any,
+                should be retried.
+            timeout (float): The timeout for this request.
 
         Returns:
             :class:`~google.cloud.firestore_v1.types.WriteResult`:
@@ -124,7 +133,8 @@ def set(self, document_data: dict, merge: bool = False) -> Any:
         """
         batch = self._client.batch()
         batch.set(self, document_data, merge=merge)
-        write_results = batch.commit()
+        kwargs = self._make_retry_timeout_kwargs(retry, timeout)
+        write_results = batch.commit(**kwargs)
         return _first_write_result(write_results)
 
     def update(self, field_updates: dict, option: _helpers.WriteOption = None) -> Any:
diff --git a/tests/unit/v1/test_document.py b/tests/unit/v1/test_document.py
index c92ce6424..1b71c40ac 100644
--- a/tests/unit/v1/test_document.py
+++ b/tests/unit/v1/test_document.py
@@ -168,7 +168,7 @@ def _write_pb_for_set(document_path, document_data, merge):
             write_pbs._pb.update_mask.CopyFrom(mask._pb)
         return write_pbs
 
-    def _set_helper(self, merge=False, **option_kwargs):
+    def _set_helper(self, merge=False, retry=None, timeout=None, **option_kwargs):
         # Create a minimal fake GAPIC with a dummy response.
         firestore_api = mock.Mock(spec=["commit"])
         firestore_api.commit.return_value = self._make_commit_repsonse()
@@ -180,7 +180,16 @@ def _set_helper(self, merge=False, **option_kwargs):
         # Actually make a document and call create().
         document = self._make_one("User", "Interface", client=client)
         document_data = {"And": 500, "Now": b"\xba\xaa\xaa \xba\xaa\xaa"}
-        write_result = document.set(document_data, merge)
+
+        kwargs = {}
+
+        if retry is not None:
+            kwargs["retry"] = retry
+
+        if timeout is not None:
+            kwargs["timeout"] = timeout
+
+        write_result = document.set(document_data, merge, **kwargs)
 
         # Verify the response and the mocks.
         self.assertIs(write_result, mock.sentinel.write_result)
@@ -193,11 +202,19 @@ def _set_helper(self, merge=False, **option_kwargs):
                 "transaction": None,
             },
             metadata=client._rpc_metadata,
+            **kwargs,
         )
 
     def test_set(self):
         self._set_helper()
 
+    def test_set_w_retry_timeout(self):
+        from google.api_core.retry import Retry
+
+        retry = Retry(predicate=object())
+        timeout = 123.0
+        self._set_helper(retry=retry, timeout=timeout)
+
     def test_set_merge(self):
         self._set_helper(merge=True)