feat: add retry/timeout to 'collection.CollectionReference.stream'

Toward #221
googleapis · Oct 13, 2020 · 16a1e34 · 16a1e34
1 parent 49d3b03
commit 16a1e34
Show file tree

Hide file tree

Showing 2 changed files with 25 additions and 2 deletions.
diff --git a/google/cloud/firestore_v1/collection.py b/google/cloud/firestore_v1/collection.py
@@ -185,7 +185,10 @@ def get(
         return query.get(transaction=transaction, **kwargs)
 
     def stream(
-        self, transaction: Transaction = None
+        self,
+        transaction: Transaction = None,
+        retry: retries.Retry = None,
+        timeout: float = None,
     ) -> Generator[document.DocumentSnapshot, Any, None]:
         """Read the documents in this collection.
 
@@ -208,13 +211,17 @@ def stream(
             transaction (Optional[:class:`~google.cloud.firestore_v1.transaction.\
                 Transaction`]):
                 An existing transaction that the query will run in.
+            retry (google.api_core.retry.Retry): Designation of what errors, if any,
+                should be retried.
+            timeout (float): The timeout for this request.
 
         Yields:
             :class:`~google.cloud.firestore_v1.document.DocumentSnapshot`:
             The next document that fulfills the query.
         """
         query = query_mod.Query(self)
-        return query.stream(transaction=transaction)
+        kwargs = self._make_retry_timeout_kwargs(retry, timeout)
+        return query.stream(transaction=transaction, **kwargs)
 
     def on_snapshot(self, callback: Callable) -> Watch:
         """Monitor the documents in this collection.

diff --git a/tests/unit/v1/test_collection.py b/tests/unit/v1/test_collection.py
@@ -326,6 +326,22 @@ def test_stream(self, query_class):
         self.assertIs(stream_response, query_instance.stream.return_value)
         query_instance.stream.assert_called_once_with(transaction=None)
 
+    @mock.patch("google.cloud.firestore_v1.query.Query", autospec=True)
+    def test_stream_w_retry_timeout(self, query_class):
+        from google.api_core.retry import Retry
+
+        retry = Retry(predicate=object())
+        timeout = 123.0
+        collection = self._make_one("collection")
+        stream_response = collection.stream(retry=retry, timeout=timeout)
+
+        query_class.assert_called_once_with(collection)
+        query_instance = query_class.return_value
+        self.assertIs(stream_response, query_instance.stream.return_value)
+        query_instance.stream.assert_called_once_with(
+            transaction=None, retry=retry, timeout=timeout,
+        )
+
     @mock.patch("google.cloud.firestore_v1.query.Query", autospec=True)
     def test_stream_with_transaction(self, query_class):
         collection = self._make_one("collection")