feat: asyncio microgen query

google/cloud/firestore_v1/async_client.py

@@ -252,7 +252,7 @@ async def get_all(self, references, field_paths=None, transaction=None):
             metadata=self._rpc_metadata,
         )
 
-        for get_doc_response in response_iterator:
+        async for get_doc_response in response_iterator:
             yield _parse_batch_get(get_doc_response, reference_map, self)
 
     async def collections(self):

google/cloud/firestore_v1/async_query.py

@@ -158,7 +158,7 @@ async def stream(self, transaction=None):
             metadata=self._client._rpc_metadata,
         )
 
-        for response in response_iterator:
+        async for response in response_iterator:
             if self._all_descendants:
                 snapshot = _collection_group_query_response_to_snapshot(
                     response, self._parent

tests/unit/v1/test__helpers.py

@@ -25,6 +25,15 @@ async def __call__(self, *args, **kwargs):
         return super(AsyncMock, self).__call__(*args, **kwargs)
 
 
+class AsyncIter:
+    def __init__(self, items):
+        self.items = items
+
+    async def __aiter__(self, **_):
+        for i in self.items:
+            yield i
+
+
 class TestGeoPoint(unittest.TestCase):
     @staticmethod
     def _get_target_class():

tests/unit/v1/test_async_client.py

@@ -18,7 +18,7 @@
 import aiounittest
 
 import mock
-from tests.unit.v1.test__helpers import AsyncMock
+from tests.unit.v1.test__helpers import AsyncMock, AsyncIter
 
 
 class TestAsyncClient(aiounittest.AsyncTestCase):
@@ -237,7 +237,7 @@ def _next_page(self):
     async 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"])
-        response_iterator = iter(document_pbs)
+        response_iterator = AsyncIter(document_pbs)
         firestore_api.batch_get_documents.return_value = response_iterator
 
         # Attach the fake GAPIC to a real client.

tests/unit/v1/test_async_collection.py

@@ -17,16 +17,7 @@
 import aiounittest
 
 import mock
-from tests.unit.v1.test__helpers import AsyncMock
-
-
-class MockAsyncIter:
-    def __init__(self, count):
-        self.count = count
-
-    async def __aiter__(self, **_):
-        for i in range(self.count):
-            yield i
+from tests.unit.v1.test__helpers import AsyncMock, AsyncIter
 
 
 class TestAsyncCollectionReference(aiounittest.AsyncTestCase):
@@ -258,7 +249,7 @@ async def test_list_documents_w_page_size(self):
     async def test_get(self, query_class):
         import warnings
 
-        query_class.return_value.stream.return_value = MockAsyncIter(3)
+        query_class.return_value.stream.return_value = AsyncIter(range(3))
 
         collection = self._make_one("collection")
         with warnings.catch_warnings(record=True) as warned:
@@ -280,7 +271,7 @@ async def test_get(self, query_class):
     async def test_get_with_transaction(self, query_class):
         import warnings
 
-        query_class.return_value.stream.return_value = MockAsyncIter(3)
+        query_class.return_value.stream.return_value = AsyncIter(range(3))
 
         collection = self._make_one("collection")
         transaction = mock.sentinel.txn
@@ -301,7 +292,7 @@ async def test_get_with_transaction(self, query_class):
     @mock.patch("google.cloud.firestore_v1.async_query.AsyncQuery", autospec=True)
     @pytest.mark.asyncio
     async def test_stream(self, query_class):
-        query_class.return_value.stream.return_value = MockAsyncIter(3)
+        query_class.return_value.stream.return_value = AsyncIter(range(3))
 
         collection = self._make_one("collection")
         stream_response = collection.stream()
@@ -316,7 +307,7 @@ async def test_stream(self, query_class):
     @mock.patch("google.cloud.firestore_v1.async_query.AsyncQuery", autospec=True)
     @pytest.mark.asyncio
     async def test_stream_with_transaction(self, query_class):
-        query_class.return_value.stream.return_value = MockAsyncIter(3)
+        query_class.return_value.stream.return_value = AsyncIter(range(3))
 
         collection = self._make_one("collection")
         transaction = mock.sentinel.txn

tests/unit/v1/test_async_query.py

@@ -17,7 +17,7 @@
 import aiounittest
 
 import mock
-
+from tests.unit.v1.test__helpers import AsyncIter
 from tests.unit.v1.test_base_query import _make_credentials, _make_query_response
 
 
@@ -59,7 +59,7 @@ async def test_get(self):
         import warnings
 
         with mock.patch.object(self._get_target_class(), "stream") as stream_mock:
-            stream_mock.return_value = MockAsyncIter(3)
+            stream_mock.return_value = AsyncIter(range(3))
 
             # Create a minimal fake GAPIC.
             firestore_api = mock.Mock(spec=["run_query"])
@@ -81,7 +81,7 @@ async def test_get(self):
                 # Verify that `get` merely wraps `stream`.
                 stream_mock.assert_called_once()
                 self.assertIsInstance(get_response, types.AsyncGeneratorType)
-                self.assertEqual(returned, list(range(stream_mock.return_value.count)))
+                self.assertEqual(returned, list(stream_mock.return_value.items))
 
             # Verify the deprecation.
             self.assertEqual(len(warned), 1)
@@ -104,7 +104,7 @@ async def test_stream_simple(self):
         name = "{}/sleep".format(expected_prefix)
         data = {"snooze": 10}
         response_pb = _make_query_response(name=name, data=data)
-        firestore_api.run_query.return_value = iter([response_pb])
+        firestore_api.run_query.return_value = AsyncIter([response_pb])
 
         # Execute the query and check the response.
         query = self._make_one(parent)
@@ -149,7 +149,7 @@ async def test_stream_with_transaction(self):
         name = "{}/burger".format(expected_prefix)
         data = {"lettuce": b"\xee\x87"}
         response_pb = _make_query_response(name=name, data=data)
-        firestore_api.run_query.return_value = iter([response_pb])
+        firestore_api.run_query.return_value = AsyncIter([response_pb])
 
         # Execute the query and check the response.
         query = self._make_one(parent)
@@ -176,7 +176,7 @@ async def test_stream_no_results(self):
         # Create a minimal fake GAPIC with a dummy response.
         firestore_api = mock.Mock(spec=["run_query"])
         empty_response = _make_query_response()
-        run_query_response = iter([empty_response])
+        run_query_response = AsyncIter([empty_response])
         firestore_api.run_query.return_value = run_query_response
 
         # Attach the fake GAPIC to a real client.
@@ -208,7 +208,7 @@ async def test_stream_second_response_in_empty_stream(self):
         firestore_api = mock.Mock(spec=["run_query"])
         empty_response1 = _make_query_response()
         empty_response2 = _make_query_response()
-        run_query_response = iter([empty_response1, empty_response2])
+        run_query_response = AsyncIter([empty_response1, empty_response2])
         firestore_api.run_query.return_value = run_query_response
 
         # Attach the fake GAPIC to a real client.
@@ -252,7 +252,7 @@ async def test_stream_with_skipped_results(self):
         name = "{}/clock".format(expected_prefix)
         data = {"noon": 12, "nested": {"bird": 10.5}}
         response_pb2 = _make_query_response(name=name, data=data)
-        firestore_api.run_query.return_value = iter([response_pb1, response_pb2])
+        firestore_api.run_query.return_value = AsyncIter([response_pb1, response_pb2])
 
         # Execute the query and check the response.
         query = self._make_one(parent)
@@ -293,7 +293,7 @@ async def test_stream_empty_after_first_response(self):
         data = {"lee": "hoop"}
         response_pb1 = _make_query_response(name=name, data=data)
         response_pb2 = _make_query_response()
-        firestore_api.run_query.return_value = iter([response_pb1, response_pb2])
+        firestore_api.run_query.return_value = AsyncIter([response_pb1, response_pb2])
 
         # Execute the query and check the response.
         query = self._make_one(parent)
@@ -335,7 +335,7 @@ async def test_stream_w_collection_group(self):
         data = {"lee": "hoop"}
         response_pb1 = _make_query_response(name=name, data=data)
         response_pb2 = _make_query_response()
-        firestore_api.run_query.return_value = iter([response_pb1, response_pb2])
+        firestore_api.run_query.return_value = AsyncIter([response_pb1, response_pb2])
 
         # Execute the query and check the response.
         query = self._make_one(parent)