1515import inspect
1616from unittest import mock
1717
18+ try :
19+ from unittest .mock import AsyncMock
20+ except ImportError :
21+ from mock import AsyncMock
1822import pytest
1923
2024from google .api_core import page_iterator_async
@@ -58,7 +62,7 @@ async def test_anext(self):
5862 )
5963
6064 async_iterator = PageAsyncIteratorImpl (None , None )
61- async_iterator ._next_page = mock . AsyncMock (side_effect = [page_1 , page_2 , None ])
65+ async_iterator ._next_page = AsyncMock (side_effect = [page_1 , page_2 , None ])
6266
6367 # Consume items and check the state of the async_iterator.
6468 assert async_iterator .num_results == 0
@@ -98,7 +102,7 @@ async def test__page_aiter_increment(self):
98102 page = page_iterator_async .Page (
99103 iterator , ("item" ,), page_iterator_async ._item_to_value_identity
100104 )
101- iterator ._next_page = mock . AsyncMock (side_effect = [page , None ])
105+ iterator ._next_page = AsyncMock (side_effect = [page , None ])
102106
103107 assert iterator .num_results == 0
104108
@@ -136,7 +140,7 @@ async def test__items_aiter(self):
136140 )
137141
138142 iterator = PageAsyncIteratorImpl (None , None )
139- iterator ._next_page = mock . AsyncMock (side_effect = [page1 , page2 , None ])
143+ iterator ._next_page = AsyncMock (side_effect = [page1 , page2 , None ])
140144
141145 items_aiter = iterator ._items_aiter ()
142146
@@ -159,7 +163,7 @@ async def test__items_aiter(self):
159163 @pytest .mark .asyncio
160164 async def test___aiter__ (self ):
161165 async_iterator = PageAsyncIteratorImpl (None , None )
162- async_iterator ._next_page = mock . AsyncMock (side_effect = [(1 , 2 ), (3 ,), None ])
166+ async_iterator ._next_page = AsyncMock (side_effect = [(1 , 2 ), (3 ,), None ])
163167
164168 assert not async_iterator ._started
165169
@@ -248,7 +252,7 @@ async def test_iterate(self):
248252 response1 = mock .Mock (items = ["a" , "b" ], next_page_token = "1" )
249253 response2 = mock .Mock (items = ["c" ], next_page_token = "2" )
250254 response3 = mock .Mock (items = ["d" ], next_page_token = "" )
251- method = mock . AsyncMock (side_effect = [response1 , response2 , response3 ])
255+ method = AsyncMock (side_effect = [response1 , response2 , response3 ])
252256 iterator = page_iterator_async .AsyncGRPCIterator (
253257 mock .sentinel .client , method , request , "items"
254258 )
@@ -271,7 +275,7 @@ async def test_iterate_with_max_results(self):
271275 response1 = mock .Mock (items = ["a" , "b" ], next_page_token = "1" )
272276 response2 = mock .Mock (items = ["c" ], next_page_token = "2" )
273277 response3 = mock .Mock (items = ["d" ], next_page_token = "" )
274- method = mock . AsyncMock (side_effect = [response1 , response2 , response3 ])
278+ method = AsyncMock (side_effect = [response1 , response2 , response3 ])
275279 iterator = page_iterator_async .AsyncGRPCIterator (
276280 mock .sentinel .client , method , request , "items" , max_results = 3
277281 )
0 commit comments