@@ -197,60 +197,6 @@ public void shouldFailWhenAskedForSingleRecordButResultContainsMore()
197
197
}
198
198
}
199
199
200
- @ Test
201
- public void shouldConsumeAsyncWhenResultContainsMultipleRecords ()
202
- {
203
- PullAllResponseHandler pullAllHandler = mock ( PullAllResponseHandler .class );
204
-
205
- Record record1 = new InternalRecord ( asList ( "key1" , "key2" , "key3" ), values ( 1 , 1 , 1 ) );
206
- Record record2 = new InternalRecord ( asList ( "key1" , "key2" , "key3" ), values ( 2 , 2 , 2 ) );
207
- Record record3 = new InternalRecord ( asList ( "key1" , "key2" , "key3" ), values ( 3 , 3 , 3 ) );
208
- when ( pullAllHandler .nextAsync () ).thenReturn ( completedFuture ( record1 ) )
209
- .thenReturn ( completedFuture ( record2 ) ).thenReturn ( completedFuture ( record3 ) )
210
- .thenReturn ( completedWithNull () );
211
-
212
- ResultSummary summary = mock ( ResultSummary .class );
213
- when ( pullAllHandler .summaryAsync () ).thenReturn ( completedFuture ( summary ) );
214
-
215
- InternalStatementResultCursor cursor = newCursor ( pullAllHandler );
216
-
217
- assertEquals ( summary , await ( cursor .consumeAsync () ) );
218
- verify ( pullAllHandler , times ( 4 ) ).nextAsync ();
219
- }
220
-
221
- @ Test
222
- public void shouldConsumeAsyncWhenResultContainsOneRecords ()
223
- {
224
- PullAllResponseHandler pullAllHandler = mock ( PullAllResponseHandler .class );
225
-
226
- Record record = new InternalRecord ( asList ( "key1" , "key2" ), values ( 1 , 1 ) );
227
- when ( pullAllHandler .nextAsync () ).thenReturn ( completedFuture ( record ) )
228
- .thenReturn ( completedWithNull () );
229
-
230
- ResultSummary summary = mock ( ResultSummary .class );
231
- when ( pullAllHandler .summaryAsync () ).thenReturn ( completedFuture ( summary ) );
232
-
233
- InternalStatementResultCursor cursor = newCursor ( pullAllHandler );
234
-
235
- assertEquals ( summary , await ( cursor .consumeAsync () ) );
236
- verify ( pullAllHandler , times ( 2 ) ).nextAsync ();
237
- }
238
-
239
- @ Test
240
- public void shouldConsumeAsyncWhenResultContainsNoRecords ()
241
- {
242
- PullAllResponseHandler pullAllHandler = mock ( PullAllResponseHandler .class );
243
- when ( pullAllHandler .nextAsync () ).thenReturn ( completedWithNull () );
244
-
245
- ResultSummary summary = mock ( ResultSummary .class );
246
- when ( pullAllHandler .summaryAsync () ).thenReturn ( completedFuture ( summary ) );
247
-
248
- InternalStatementResultCursor cursor = newCursor ( pullAllHandler );
249
-
250
- assertEquals ( summary , await ( cursor .consumeAsync () ) );
251
- verify ( pullAllHandler ).nextAsync ();
252
- }
253
-
254
200
@ Test
255
201
public void shouldForEachAsyncWhenResultContainsMultipleRecords ()
256
202
{
@@ -455,6 +401,38 @@ public void shouldPropagateFailureFromListAsyncWithMapFunction()
455
401
verify ( pullAllHandler ).listAsync ( mapFunction );
456
402
}
457
403
404
+ @ Test
405
+ public void shouldConsumeAsync ()
406
+ {
407
+ PullAllResponseHandler pullAllHandler = mock ( PullAllResponseHandler .class );
408
+ ResultSummary summary = mock ( ResultSummary .class );
409
+ when ( pullAllHandler .consumeAsync () ).thenReturn ( completedFuture ( summary ) );
410
+
411
+ InternalStatementResultCursor cursor = newCursor ( pullAllHandler );
412
+
413
+ assertEquals ( summary , await ( cursor .consumeAsync () ) );
414
+ }
415
+
416
+ @ Test
417
+ public void shouldPropagateFailureInConsumeAsync ()
418
+ {
419
+ PullAllResponseHandler pullAllHandler = mock ( PullAllResponseHandler .class );
420
+ RuntimeException error = new RuntimeException ( "Hi" );
421
+ when ( pullAllHandler .consumeAsync () ).thenReturn ( failedFuture ( error ) );
422
+
423
+ InternalStatementResultCursor cursor = newCursor ( pullAllHandler );
424
+
425
+ try
426
+ {
427
+ await ( cursor .consumeAsync () );
428
+ fail ( "Exception expected" );
429
+ }
430
+ catch ( RuntimeException e )
431
+ {
432
+ assertEquals ( error , e );
433
+ }
434
+ }
435
+
458
436
private static InternalStatementResultCursor newCursor ( PullAllResponseHandler pullAllHandler )
459
437
{
460
438
return new InternalStatementResultCursor ( new RunResponseHandler ( new CompletableFuture <>() ), pullAllHandler );
0 commit comments