diff --git a/gcloud-java-datastore/pom.xml b/gcloud-java-datastore/pom.xml
index a76e4f9d9cbd..5866500ba0d6 100644
--- a/gcloud-java-datastore/pom.xml
+++ b/gcloud-java-datastore/pom.xml
@@ -24,7 +24,7 @@
     <dependency>
       <groupId>com.google.apis</groupId>
       <artifactId>google-api-services-datastore-protobuf</artifactId>
-      <version>v1beta2-rev1-2.1.2</version>
+      <version>v1beta2-rev1-4.0.0</version>
       <scope>compile</scope>
       <exclusions>
         <exclusion>
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
index e9eed027e8e0..19f53f525b1a 100644
--- 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
@@ -31,6 +31,7 @@
 import com.google.api.services.datastore.DatastoreV1.RunQueryRequest;
 import com.google.api.services.datastore.DatastoreV1.RunQueryResponse;
 import com.google.common.collect.Iterators;
+import com.google.common.collect.Lists;
 import com.google.gcloud.RetryParams;
 import com.google.gcloud.datastore.Query.ResultType;
 import com.google.gcloud.datastore.StructuredQuery.OrderBy;
@@ -89,8 +90,8 @@ public class DatastoreTest {
       FullEntity.builder(INCOMPLETE_KEY2).set("str", STR_VALUE).set("bool", BOOL_VALUE)
           .set("list", LIST_VALUE1).build();
   private static final FullEntity<IncompleteKey> PARTIAL_ENTITY2 =
-      FullEntity.builder(PARTIAL_ENTITY1).remove("str").set("bool", true).
-          set("list", LIST_VALUE1.get()).build();
+      FullEntity.builder(PARTIAL_ENTITY1).remove("str").set("bool", true)
+          .set("list", LIST_VALUE1.get()).build();
   private static final FullEntity<IncompleteKey> PARTIAL_ENTITY3 =
       FullEntity.builder(PARTIAL_ENTITY1).key(IncompleteKey.builder(PROJECT_ID, KIND3).build())
           .build();
@@ -471,9 +472,13 @@ public void testQueryPaginationWithLimit() throws DatastoreException {
     EasyMock.expect(rpcFactoryMock.create(EasyMock.anyObject(DatastoreOptions.class)))
         .andReturn(rpcMock);
     List<RunQueryResponse> responses = buildResponsesForQueryPaginationWithLimit();
-    for (int i = 0; i < responses.size(); i++) {
+    List<ByteString> endCursors = Lists.newArrayListWithCapacity(responses.size());
+    for (RunQueryResponse response: responses) {
       EasyMock.expect(rpcMock.runQuery(EasyMock.anyObject(RunQueryRequest.class)))
-          .andReturn(responses.get(i));
+          .andReturn(response);
+      if (response.getBatch().getMoreResults() != QueryResultBatch.MoreResultsType.NOT_FINISHED) {
+        endCursors.add(response.getBatch().getEndCursor());
+      }
     }
     EasyMock.replay(rpcFactoryMock, rpcMock);
     Datastore mockDatastore = options.toBuilder()
@@ -483,6 +488,7 @@ public void testQueryPaginationWithLimit() throws DatastoreException {
         .service();
     int limit = 2;
     int totalCount = 0;
+    Iterator<ByteString> cursorIter = endCursors.iterator();
     StructuredQuery<Entity> query = Query.entityQueryBuilder().limit(limit).build();
     while (true) {
       QueryResults<Entity> results = mockDatastore.run(query);
@@ -492,6 +498,9 @@ public void testQueryPaginationWithLimit() throws DatastoreException {
         resultCount++;
         totalCount++;
       }
+      assertTrue(cursorIter.hasNext());
+      Cursor expectedEndCursor = Cursor.copyFrom(cursorIter.next().toByteArray());
+      assertEquals(expectedEndCursor, results.cursorAfter());
       if (resultCount < limit) {
         break;
       }
@@ -505,19 +514,20 @@ private List<RunQueryResponse> buildResponsesForQueryPaginationWithLimit() {
     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);
+    DatastoreRpc datastoreRpc = datastore.options().rpc();
     List<RunQueryResponse> responses = new ArrayList<>();
     Query<Entity> query = Query.entityQueryBuilder().build();
     RunQueryRequest.Builder requestPb = RunQueryRequest.newBuilder();
     query.populatePb(requestPb);
     QueryResultBatch queryResultBatchPb = RunQueryResponse.newBuilder()
-        .mergeFrom(((DatastoreImpl) datastore).runQuery(requestPb.build()))
+        .mergeFrom(datastoreRpc.runQuery(requestPb.build()))
         .getBatch();
     QueryResultBatch queryResultBatchPb1 = QueryResultBatch.newBuilder()
         .mergeFrom(queryResultBatchPb)
         .setMoreResults(QueryResultBatch.MoreResultsType.NOT_FINISHED)
         .clearEntityResult()
         .addAllEntityResult(queryResultBatchPb.getEntityResultList().subList(0, 1))
-        .setEndCursor(queryResultBatchPb.getEntityResultList().get(0).getCursor())
+        .setEndCursor(ByteString.copyFromUtf8("a"))
         .build();
     responses.add(RunQueryResponse.newBuilder().setBatch(queryResultBatchPb1).build());
     QueryResultBatch queryResultBatchPb2 = QueryResultBatch.newBuilder()
@@ -534,7 +544,7 @@ private List<RunQueryResponse> buildResponsesForQueryPaginationWithLimit() {
         .setMoreResults(QueryResultBatch.MoreResultsType.MORE_RESULTS_AFTER_LIMIT)
         .clearEntityResult()
         .addAllEntityResult(queryResultBatchPb.getEntityResultList().subList(2, 4))
-        .setEndCursor(queryResultBatchPb.getEntityResultList().get(3).getCursor())
+        .setEndCursor(ByteString.copyFromUtf8("b"))
         .build();
     responses.add(RunQueryResponse.newBuilder().setBatch(queryResultBatchPb3).build());
     QueryResultBatch queryResultBatchPb4 = QueryResultBatch.newBuilder()
@@ -542,7 +552,7 @@ private List<RunQueryResponse> buildResponsesForQueryPaginationWithLimit() {
         .setMoreResults(QueryResultBatch.MoreResultsType.NO_MORE_RESULTS)
         .clearEntityResult()
         .addAllEntityResult(queryResultBatchPb.getEntityResultList().subList(4, 5))
-        .setEndCursor(queryResultBatchPb.getEntityResultList().get(4).getCursor())
+        .setEndCursor(ByteString.copyFromUtf8("c"))
         .build();
     responses.add(RunQueryResponse.newBuilder().setBatch(queryResultBatchPb4).build());
     return responses;