3131import com .google .api .services .datastore .DatastoreV1 .RunQueryRequest ;
3232import com .google .api .services .datastore .DatastoreV1 .RunQueryResponse ;
3333import com .google .common .collect .Iterators ;
34+ import com .google .common .collect .Lists ;
3435import com .google .gcloud .RetryParams ;
3536import com .google .gcloud .datastore .Query .ResultType ;
3637import com .google .gcloud .datastore .StructuredQuery .OrderBy ;
@@ -89,8 +90,8 @@ public class DatastoreTest {
8990 FullEntity .builder (INCOMPLETE_KEY2 ).set ("str" , STR_VALUE ).set ("bool" , BOOL_VALUE )
9091 .set ("list" , LIST_VALUE1 ).build ();
9192 private static final FullEntity <IncompleteKey > PARTIAL_ENTITY2 =
92- FullEntity .builder (PARTIAL_ENTITY1 ).remove ("str" ).set ("bool" , true ).
93- set ("list" , LIST_VALUE1 .get ()).build ();
93+ FullEntity .builder (PARTIAL_ENTITY1 ).remove ("str" ).set ("bool" , true )
94+ . set ("list" , LIST_VALUE1 .get ()).build ();
9495 private static final FullEntity <IncompleteKey > PARTIAL_ENTITY3 =
9596 FullEntity .builder (PARTIAL_ENTITY1 ).key (IncompleteKey .builder (PROJECT_ID , KIND3 ).build ())
9697 .build ();
@@ -471,9 +472,13 @@ public void testQueryPaginationWithLimit() throws DatastoreException {
471472 EasyMock .expect (rpcFactoryMock .create (EasyMock .anyObject (DatastoreOptions .class )))
472473 .andReturn (rpcMock );
473474 List <RunQueryResponse > responses = buildResponsesForQueryPaginationWithLimit ();
474- for (int i = 0 ; i < responses .size (); i ++) {
475+ List <ByteString > endCursors = Lists .newArrayListWithCapacity (responses .size ());
476+ for (RunQueryResponse response : responses ) {
475477 EasyMock .expect (rpcMock .runQuery (EasyMock .anyObject (RunQueryRequest .class )))
476- .andReturn (responses .get (i ));
478+ .andReturn (response );
479+ if (response .getBatch ().getMoreResults () != QueryResultBatch .MoreResultsType .NOT_FINISHED ) {
480+ endCursors .add (response .getBatch ().getEndCursor ());
481+ }
477482 }
478483 EasyMock .replay (rpcFactoryMock , rpcMock );
479484 Datastore mockDatastore = options .toBuilder ()
@@ -483,6 +488,7 @@ public void testQueryPaginationWithLimit() throws DatastoreException {
483488 .service ();
484489 int limit = 2 ;
485490 int totalCount = 0 ;
491+ Iterator <ByteString > cursorIter = endCursors .iterator ();
486492 StructuredQuery <Entity > query = Query .entityQueryBuilder ().limit (limit ).build ();
487493 while (true ) {
488494 QueryResults <Entity > results = mockDatastore .run (query );
@@ -492,6 +498,9 @@ public void testQueryPaginationWithLimit() throws DatastoreException {
492498 resultCount ++;
493499 totalCount ++;
494500 }
501+ assertTrue (cursorIter .hasNext ());
502+ Cursor expectedEndCursor = Cursor .copyFrom (cursorIter .next ().toByteArray ());
503+ assertEquals (expectedEndCursor , results .cursorAfter ());
495504 if (resultCount < limit ) {
496505 break ;
497506 }
@@ -505,19 +514,20 @@ private List<RunQueryResponse> buildResponsesForQueryPaginationWithLimit() {
505514 Entity entity4 = Entity .builder (KEY4 ).set ("value" , StringValue .of ("value" )).build ();
506515 Entity entity5 = Entity .builder (KEY5 ).set ("value" , "value" ).build ();
507516 datastore .add (ENTITY3 , entity4 , entity5 );
517+ DatastoreRpc datastoreRpc = datastore .options ().rpc ();
508518 List <RunQueryResponse > responses = new ArrayList <>();
509519 Query <Entity > query = Query .entityQueryBuilder ().build ();
510520 RunQueryRequest .Builder requestPb = RunQueryRequest .newBuilder ();
511521 query .populatePb (requestPb );
512522 QueryResultBatch queryResultBatchPb = RunQueryResponse .newBuilder ()
513- .mergeFrom ((( DatastoreImpl ) datastore ) .runQuery (requestPb .build ()))
523+ .mergeFrom (datastoreRpc .runQuery (requestPb .build ()))
514524 .getBatch ();
515525 QueryResultBatch queryResultBatchPb1 = QueryResultBatch .newBuilder ()
516526 .mergeFrom (queryResultBatchPb )
517527 .setMoreResults (QueryResultBatch .MoreResultsType .NOT_FINISHED )
518528 .clearEntityResult ()
519529 .addAllEntityResult (queryResultBatchPb .getEntityResultList ().subList (0 , 1 ))
520- .setEndCursor (queryResultBatchPb . getEntityResultList (). get ( 0 ). getCursor ( ))
530+ .setEndCursor (ByteString . copyFromUtf8 ( "a" ))
521531 .build ();
522532 responses .add (RunQueryResponse .newBuilder ().setBatch (queryResultBatchPb1 ).build ());
523533 QueryResultBatch queryResultBatchPb2 = QueryResultBatch .newBuilder ()
@@ -534,15 +544,15 @@ private List<RunQueryResponse> buildResponsesForQueryPaginationWithLimit() {
534544 .setMoreResults (QueryResultBatch .MoreResultsType .MORE_RESULTS_AFTER_LIMIT )
535545 .clearEntityResult ()
536546 .addAllEntityResult (queryResultBatchPb .getEntityResultList ().subList (2 , 4 ))
537- .setEndCursor (queryResultBatchPb . getEntityResultList (). get ( 3 ). getCursor ( ))
547+ .setEndCursor (ByteString . copyFromUtf8 ( "b" ))
538548 .build ();
539549 responses .add (RunQueryResponse .newBuilder ().setBatch (queryResultBatchPb3 ).build ());
540550 QueryResultBatch queryResultBatchPb4 = QueryResultBatch .newBuilder ()
541551 .mergeFrom (queryResultBatchPb )
542552 .setMoreResults (QueryResultBatch .MoreResultsType .NO_MORE_RESULTS )
543553 .clearEntityResult ()
544554 .addAllEntityResult (queryResultBatchPb .getEntityResultList ().subList (4 , 5 ))
545- .setEndCursor (queryResultBatchPb . getEntityResultList (). get ( 4 ). getCursor ( ))
555+ .setEndCursor (ByteString . copyFromUtf8 ( "c" ))
546556 .build ();
547557 responses .add (RunQueryResponse .newBuilder ().setBatch (queryResultBatchPb4 ).build ());
548558 return responses ;
0 commit comments