diff --git a/google/cloud/firestore_v1/async_query.py b/google/cloud/firestore_v1/async_query.py index dea0c960b..a4a46d6ec 100644 --- a/google/cloud/firestore_v1/async_query.py +++ b/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 diff --git a/tests/unit/v1/test_async_query.py b/tests/unit/v1/test_async_query.py index 289564606..1bbbf9ff7 100644 --- a/tests/unit/v1/test_async_query.py +++ b/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)