@@ -198,6 +198,15 @@ async def fetch_and_compare_all_records(
198
198
if limit is None :
199
199
assert result ._closed
200
200
elif method == "next" :
201
+ n = len (expected_records ) if limit is None else limit
202
+ for _ in range (n ):
203
+ record = await AsyncUtil .next (result )
204
+ received_records .append ([record .get (key , None )])
205
+ if limit is None :
206
+ with pytest .raises (StopAsyncIteration ):
207
+ await AsyncUtil .next (result )
208
+ assert result ._closed
209
+ elif method == "one iter" :
201
210
iter_ = AsyncUtil .iter (result )
202
211
n = len (expected_records ) if limit is None else limit
203
212
for _ in range (n ):
@@ -223,7 +232,8 @@ async def fetch_and_compare_all_records(
223
232
assert received_records == expected_records
224
233
225
234
226
- @pytest .mark .parametrize ("method" , ("for loop" , "next" , "new iter" ))
235
+ @pytest .mark .parametrize ("method" ,
236
+ ("for loop" , "next" , "one iter" , "new iter" ))
227
237
@pytest .mark .parametrize ("records" , (
228
238
[],
229
239
[[42 ]],
@@ -237,7 +247,8 @@ async def test_result_iteration(method, records):
237
247
await fetch_and_compare_all_records (result , "x" , records , method )
238
248
239
249
240
- @pytest .mark .parametrize ("method" , ("for loop" , "next" , "new iter" ))
250
+ @pytest .mark .parametrize ("method" ,
251
+ ("for loop" , "next" , "one iter" , "new iter" ))
241
252
@pytest .mark .parametrize ("invert_fetch" , (True , False ))
242
253
@mark_async_test
243
254
async def test_parallel_result_iteration (method , invert_fetch ):
@@ -266,7 +277,8 @@ async def test_parallel_result_iteration(method, invert_fetch):
266
277
)
267
278
268
279
269
- @pytest .mark .parametrize ("method" , ("for loop" , "next" , "new iter" ))
280
+ @pytest .mark .parametrize ("method" ,
281
+ ("for loop" , "next" , "one iter" , "new iter" ))
270
282
@pytest .mark .parametrize ("invert_fetch" , (True , False ))
271
283
@mark_async_test
272
284
async def test_interwoven_result_iteration (method , invert_fetch ):
0 commit comments