feat: add retry/timeout to 'client.Client.collections'

Toward #221

google/cloud/firestore_v1/client.py

@@ -279,16 +279,26 @@ def get_all(
         for get_doc_response in response_iterator:
             yield _parse_batch_get(get_doc_response, reference_map, self)
 
-    def collections(self) -> Generator[Any, Any, None]:
+    def collections(
+        self, retry: retries.Retry = None, timeout: float = None,
+    ) -> Generator[Any, Any, None]:
         """List top-level collections of the client's database.
 
+        Args:
+            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.
         """
+        kwargs = self._make_retry_timeout_kwargs(retry, timeout)
+
         iterator = self._firestore_api.list_collection_ids(
             request={"parent": "{}/documents".format(self._database_string)},
             metadata=self._rpc_metadata,
+            **kwargs,
         )
 
         while True:

tests/unit/v1/test_client.py

@@ -232,6 +232,51 @@ def _next_page(self):
             request={"parent": base_path}, metadata=client._rpc_metadata
         )
 
+    def test_collections_w_retry_timeout(self):
+        from google.api_core.page_iterator import Iterator
+        from google.api_core.page_iterator import Page
+        from google.api_core.retry import Retry
+        from google.cloud.firestore_v1.collection import CollectionReference
+
+        collection_ids = ["users", "projects"]
+        retry = Retry(predicate=object())
+        timeout = 123.0
+        client = self._make_default_one()
+        firestore_api = mock.Mock(spec=["list_collection_ids"])
+        client._firestore_api_internal = firestore_api
+
+        # TODO(microgen): list_collection_ids isn't a pager.
+        # https://github.com/googleapis/gapic-generator-python/issues/516
+        class _Iterator(Iterator):
+            def __init__(self, pages):
+                super(_Iterator, self).__init__(client=None)
+                self._pages = pages
+                self.collection_ids = pages[0]
+
+            def _next_page(self):
+                if self._pages:
+                    page, self._pages = self._pages[0], self._pages[1:]
+                    return Page(self, page, self.item_to_value)
+
+        iterator = _Iterator(pages=[collection_ids])
+        firestore_api.list_collection_ids.return_value = iterator
+
+        collections = list(client.collections(retry=retry, timeout=timeout))
+
+        self.assertEqual(len(collections), len(collection_ids))
+        for collection, collection_id in zip(collections, collection_ids):
+            self.assertIsInstance(collection, CollectionReference)
+            self.assertEqual(collection.parent, None)
+            self.assertEqual(collection.id, collection_id)
+
+        base_path = client._database_string + "/documents"
+        firestore_api.list_collection_ids.assert_called_once_with(
+            request={"parent": base_path},
+            retry=retry,
+            timeout=timeout,
+            metadata=client._rpc_metadata,
+        )
+
     def _get_all_helper(self, client, references, document_pbs, **kwargs):
         # Create a minimal fake GAPIC with a dummy response.
         firestore_api = mock.Mock(spec=["batch_get_documents"])