feat: add retry/timeout to 'document.Document.collections'

Toward #221

google/cloud/firestore_v1/document.py

@@ -372,23 +372,31 @@ def get(
             update_time=update_time,
         )
 
-    def collections(self, page_size: int = None) -> Generator[Any, Any, None]:
+    def collections(
+        self, page_size: int = None, retry: retries.Retry = None, timeout: float = None,
+    ) -> Generator[Any, Any, None]:
         """List subcollections of the current document.
 
         Args:
             page_size (Optional[int]]): The maximum number of collections
-            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.CollectionReference`]:
                 iterator of subcollections of the current document. If the
                 document does not exist at the time of `snapshot`, the
                 iterator will be empty
         """
+        kwargs = self._make_retry_timeout_kwargs(retry, timeout)
+
         iterator = self._client._firestore_api.list_collection_ids(
             request={"parent": self._document_path, "page_size": page_size},
             metadata=self._client._rpc_metadata,
+            **kwargs,
         )
 
         while True:

tests/unit/v1/test_document.py

@@ -395,7 +395,7 @@ def _get_helper(
                 "transaction": expected_transaction_id,
             },
             metadata=client._rpc_metadata,
-            **kwargs
+            **kwargs,
         )
 
     def test_get_not_found(self):
@@ -424,7 +424,7 @@ def test_get_with_multiple_field_paths(self):
     def test_get_with_transaction(self):
         self._get_helper(use_transaction=True)
 
-    def _collections_helper(self, page_size=None):
+    def _collections_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.collection import CollectionReference
@@ -450,12 +450,20 @@ def _next_page(self):
         client = _make_client()
         client._firestore_api_internal = api_client
 
+        kwargs = {}
+
+        if retry is not None:
+            kwargs["retry"] = retry
+
+        if timeout is not None:
+            kwargs["timeout"] = timeout
+
         # Actually make a document and call delete().
         document = self._make_one("where", "we-are", client=client)
         if page_size is not None:
-            collections = list(document.collections(page_size=page_size))
+            collections = list(document.collections(page_size=page_size, **kwargs))
         else:
-            collections = list(document.collections())
+            collections = list(document.collections(**kwargs))
 
         # Verify the response and the mocks.
         self.assertEqual(len(collections), len(collection_ids))
@@ -467,6 +475,7 @@ def _next_page(self):
         api_client.list_collection_ids.assert_called_once_with(
             request={"parent": document._document_path, "page_size": page_size},
             metadata=client._rpc_metadata,
+            **kwargs,
         )
 
     def test_collections_wo_page_size(self):
@@ -475,6 +484,13 @@ def test_collections_wo_page_size(self):
     def test_collections_w_page_size(self):
         self._collections_helper(page_size=10)
 
+    def test_collections_w_retry_timeout(self):
+        from google.api_core.retry import Retry
+
+        retry = Retry(predicate=object())
+        timeout = 123.0
+        self._collections_helper(retry=retry, timeout=timeout)
+
     @mock.patch("google.cloud.firestore_v1.document.Watch", autospec=True)
     def test_on_snapshot(self, watch):
         client = mock.Mock(_database_string="sprinklez", spec=["_database_string"])