feat: add retry/timeout to 'collection.CollectionReference.list_docum…

…ents'

Toward #221

google/cloud/firestore_v1/collection.py

@@ -120,13 +120,18 @@ def add(
         write_result = document_ref.create(document_data, **kwargs)
         return write_result.update_time, document_ref
 
-    def list_documents(self, page_size: int = None) -> Generator[Any, Any, None]:
+    def list_documents(
+        self, page_size: int = None, retry: retries.Retry = None, timeout: float = None,
+    ) -> Generator[Any, Any, None]:
         """List all subdocuments of the current collection.
 
         Args:
             page_size (Optional[int]]): The maximum number of documents
-            in each page of results from this request. Non-positive values
-            are ignored. Defaults to a sensible value set by the API.
+                in each page of results from this request. Non-positive values
+                are ignored. Defaults to a sensible value set by the API.
+            retry (google.api_core.retry.Retry): Designation of what errors, if any,
+                should be retried.
+            timeout (float): The timeout for this request.
 
         Returns:
             Sequence[:class:`~google.cloud.firestore_v1.collection.DocumentReference`]:
@@ -135,6 +140,7 @@ def list_documents(self, page_size: int = None) -> Generator[Any, Any, None]:
                 iterator will be empty
         """
         parent, _ = self._parent_info()
+        kwargs = self._make_retry_timeout_kwargs(retry, timeout)
 
         iterator = self._client._firestore_api.list_documents(
             request={
@@ -144,6 +150,7 @@ def list_documents(self, page_size: int = None) -> Generator[Any, Any, None]:
                 "show_missing": True,
             },
             metadata=self._client._rpc_metadata,
+            **kwargs,
         )
         return (_item_to_document_ref(self, i) for i in iterator)
 

tests/unit/v1/test_collection.py

@@ -201,7 +201,7 @@ def test_add_w_retry_timeout(self):
         timeout = 123.0
         self._add_helper(retry=retry, timeout=timeout)
 
-    def _list_documents_helper(self, page_size=None):
+    def _list_documents_helper(self, page_size=None, retry=None, timeout=None):
         from google.api_core.page_iterator import Iterator
         from google.api_core.page_iterator import Page
         from google.cloud.firestore_v1.document import DocumentReference
@@ -230,10 +230,18 @@ def _next_page(self):
         client._firestore_api_internal = api_client
         collection = self._make_one("collection", client=client)
 
+        kwargs = {}
+
+        if retry is not None:
+            kwargs["retry"] = retry
+
+        if timeout is not None:
+            kwargs["timeout"] = timeout
+
         if page_size is not None:
-            documents = list(collection.list_documents(page_size=page_size))
+            documents = list(collection.list_documents(page_size=page_size, **kwargs))
         else:
-            documents = list(collection.list_documents())
+            documents = list(collection.list_documents(**kwargs))
 
         # Verify the response and the mocks.
         self.assertEqual(len(documents), len(document_ids))
@@ -251,11 +259,19 @@ def _next_page(self):
                 "show_missing": True,
             },
             metadata=client._rpc_metadata,
+            **kwargs,
         )
 
     def test_list_documents_wo_page_size(self):
         self._list_documents_helper()
 
+    def test_list_documents_w_retry_timeout(self):
+        from google.api_core.retry import Retry
+
+        retry = Retry(predicate=object())
+        timeout = 123.0
+        self._list_documents_helper(retry=retry, timeout=timeout)
+
     def test_list_documents_w_page_size(self):
         self._list_documents_helper(page_size=25)