1717import aiounittest
1818
1919import mock
20+ from tests .unit .v1 .test__helpers import AsyncMock
2021
2122
2223class MockAsyncIter :
@@ -196,7 +197,6 @@ async def _list_documents_helper(self, page_size=None):
196197 from google .api_core .page_iterator import Iterator
197198 from google .api_core .page_iterator import Page
198199 from google .cloud .firestore_v1 .async_document import AsyncDocumentReference
199- from google .cloud .firestore_v1 .services .firestore .client import FirestoreClient
200200 from google .cloud .firestore_v1 .types .document import Document
201201
202202 class _Iterator (Iterator ):
@@ -216,9 +216,10 @@ def _next_page(self):
216216 Document (name = template .format (document_id )) for document_id in document_ids
217217 ]
218218 iterator = _Iterator (pages = [documents ])
219- api_client = mock .create_autospec (FirestoreClient )
220- api_client .list_documents .return_value = iterator
221- client ._firestore_api_internal = api_client
219+ firestore_api = AsyncMock ()
220+ firestore_api .mock_add_spec (spec = ["list_documents" ])
221+ firestore_api .list_documents .return_value = iterator
222+ client ._firestore_api_internal = firestore_api
222223 collection = self ._make_one ("collection" , client = client )
223224
224225 if page_size is not None :
@@ -234,7 +235,7 @@ def _next_page(self):
234235 self .assertEqual (document .id , document_id )
235236
236237 parent , _ = collection ._parent_info ()
237- api_client .list_documents .assert_called_once_with (
238+ firestore_api .list_documents .assert_called_once_with (
238239 request = {
239240 "parent" : parent ,
240241 "collection_id" : collection .id ,
0 commit comments