Set namespace in Gql's nextQuery and add another batch to tests

Ajay Kannan · Ajay Kannan · commit b270c64f96c5 · 2015-10-21T13:08:18.000-07:00
diff --git a/gcloud-java-datastore/src/main/java/com/google/gcloud/datastore/GqlQuery.java b/gcloud-java-datastore/src/main/java/com/google/gcloud/datastore/GqlQuery.java
@@ -381,6 +381,7 @@ protected void populatePb(com.google.datastore.v1beta3.RunQueryRequest.Builder r
   @Override
   protected StructuredQuery<V> nextQuery(com.google.datastore.v1beta3.RunQueryResponse responsePb) {
     return new StructuredQuery.Builder<>(type())
+        .namespace(namespace())
         .mergeFrom(responsePb.getQuery())
         .prepareNext(responsePb.getBatch())
         .build();
diff --git a/gcloud-java-datastore/src/test/java/com/google/gcloud/datastore/DatastoreTest.java b/gcloud-java-datastore/src/test/java/com/google/gcloud/datastore/DatastoreTest.java
@@ -406,22 +406,18 @@ public void testRunGqlQueryWithCasting() {
   }
 
   @Test
-  public void testGqlQueryPagination() {
+  public void testGqlQueryPagination() throws DatastoreRpcException {
     DatastoreRpcFactory rpcFactoryMock = EasyMock.createStrictMock(DatastoreRpcFactory.class);
     DatastoreRpc rpcMock = EasyMock.createStrictMock(DatastoreRpc.class);
     EasyMock.expect(rpcFactoryMock.create(EasyMock.anyObject(DatastoreOptions.class)))
         .andReturn(rpcMock);
     List<com.google.datastore.v1beta3.RunQueryResponse> responses =
         buildResponsesForQueryPagination();
     for (int i = 0; i < responses.size(); i++) {
-      try {
         EasyMock
             .expect(rpcMock.runQuery(
                 EasyMock.anyObject(com.google.datastore.v1beta3.RunQueryRequest.class)))
             .andReturn(responses.get(i));
-      } catch (DatastoreRpcException e) {
-        fail("Unexpected DatastoreRpcException");
-      }
     }
     EasyMock.replay(rpcFactoryMock, rpcMock);
     DatastoreOptions options =
@@ -437,7 +433,7 @@ public void testGqlQueryPagination() {
       count += 1;
       results.next();
     }
-    assertEquals(count, responses.size());
+    assertEquals(count, 5);
     EasyMock.verify(rpcFactoryMock, rpcMock);
   }
 
@@ -486,22 +482,18 @@ public void testRunStructuredQuery() {
   }
 
   @Test
-  public void testStructuredQueryPagination() {
+  public void testStructuredQueryPagination() throws DatastoreRpcException {
     DatastoreRpcFactory rpcFactoryMock = EasyMock.createStrictMock(DatastoreRpcFactory.class);
     DatastoreRpc rpcMock = EasyMock.createStrictMock(DatastoreRpc.class);
     EasyMock.expect(rpcFactoryMock.create(EasyMock.anyObject(DatastoreOptions.class)))
         .andReturn(rpcMock);
     List<com.google.datastore.v1beta3.RunQueryResponse> responses =
         buildResponsesForQueryPagination();
     for (int i = 0; i < responses.size(); i++) {
-      try {
         EasyMock
             .expect(rpcMock.runQuery(
                 EasyMock.anyObject(com.google.datastore.v1beta3.RunQueryRequest.class)))
             .andReturn(responses.get(i));
-      } catch (DatastoreRpcException e) {
-        fail("Unexpected DatastoreRpcException");
-      }
     }
     EasyMock.replay(rpcFactoryMock, rpcMock);
     DatastoreOptions options =
@@ -516,11 +508,14 @@ public void testStructuredQueryPagination() {
       count += 1;
       results.next();
     }
-    assertEquals(count, responses.size());
+    assertEquals(count, 5);
     EasyMock.verify(rpcFactoryMock, rpcMock);
   }
 
   private List<com.google.datastore.v1beta3.RunQueryResponse> buildResponsesForQueryPagination() {
+    Entity entity4 = Entity.builder(KEY4).set("value", StringValue.of("value")).build();
+    Entity entity5 = Entity.builder(KEY5).set("value", "value").build();
+    datastore.add(ENTITY3, entity4, entity5);
     List<com.google.datastore.v1beta3.RunQueryResponse> responses = new ArrayList<>();
     Query<Key> query = Query.keyQueryBuilder().build();
     com.google.datastore.v1beta3.RunQueryRequest.Builder requestPb =
@@ -547,15 +542,28 @@ private List<com.google.datastore.v1beta3.RunQueryResponse> buildResponsesForQue
         com.google.datastore.v1beta3.QueryResultBatch.newBuilder()
             .mergeFrom(queryResultBatchPb)
             .setMoreResults(
-                com.google.datastore.v1beta3.QueryResultBatch.MoreResultsType.NO_MORE_RESULTS)
+                com.google.datastore.v1beta3.QueryResultBatch.MoreResultsType.NOT_FINISHED)
             .clearEntityResults()
-            .addAllEntityResults(queryResultBatchPb.getEntityResultsList().subList(1, 2))
-            .setEndCursor(queryResultBatchPb.getEntityResultsList().get(1).getCursor())
+            .addAllEntityResults(queryResultBatchPb.getEntityResultsList().subList(1, 3))
+            .setEndCursor(queryResultBatchPb.getEntityResultsList().get(2).getCursor())
             .build();
     responses.add(
         com.google.datastore.v1beta3.RunQueryResponse.newBuilder()
             .setBatch(queryResultBatchPb2)
             .build());
+    com.google.datastore.v1beta3.QueryResultBatch queryResultBatchPb3 =
+        com.google.datastore.v1beta3.QueryResultBatch.newBuilder()
+            .mergeFrom(queryResultBatchPb)
+            .setMoreResults(
+                com.google.datastore.v1beta3.QueryResultBatch.MoreResultsType.NO_MORE_RESULTS)
+            .clearEntityResults()
+            .addAllEntityResults(queryResultBatchPb.getEntityResultsList().subList(3, 5))
+            .setEndCursor(queryResultBatchPb.getEntityResultsList().get(4).getCursor())
+            .build();
+    responses.add(
+        com.google.datastore.v1beta3.RunQueryResponse.newBuilder()
+            .setBatch(queryResultBatchPb3)
+            .build());
     return responses;
   }
 
