Merge pull request #285 from ajkannan/test-deferred

Add test for deferred results in datastore get
googleapis · Oct 23, 2015 · 36af448 · 36af448
2 parents e4da160 + d4530f6
commit 36af448
Showing 1 changed file with 84 additions and 2 deletions.
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
@@ -35,6 +35,7 @@
 import com.google.gcloud.datastore.StructuredQuery.PropertyFilter;
 import com.google.gcloud.datastore.testing.LocalGcdHelper;
 import com.google.gcloud.spi.DatastoreRpc;
+import com.google.gcloud.spi.DatastoreRpc.DatastoreRpcException;
 import com.google.gcloud.spi.DatastoreRpc.DatastoreRpcException.Reason;
 import com.google.gcloud.spi.DatastoreRpcFactory;
 
@@ -49,9 +50,12 @@
 import org.junit.runners.JUnit4;
 
 import java.io.IOException;
+import java.util.ArrayList;
 import java.util.Collections;
+import java.util.HashSet;
 import java.util.Iterator;
 import java.util.List;
+import java.util.Set;
 
 @RunWith(JUnit4.class)
 public class DatastoreTest {
@@ -520,7 +524,7 @@ public void testGet() {
   }
 
   @Test
-  public void testGetArray() {
+  public void testGetArrayNoDeferredResults() {
     datastore.put(ENTITY3);
     Iterator<Entity> result =
         datastore.fetch(KEY1, Key.builder(KEY1).name("bla").build(), KEY2, KEY3).iterator();
@@ -546,7 +550,85 @@ public void testGetArray() {
       // expected - no such property
     }
     assertFalse(result.hasNext());
-    // TODO(ozarov): construct a test to verify more results
+  }
+
+  public void testGetArrayDeferredResults() throws DatastoreRpcException {
+    Set<Key> requestedKeys = new HashSet<>();
+    requestedKeys.add(KEY1);
+    requestedKeys.add(KEY2);
+    requestedKeys.add(KEY3);
+    requestedKeys.add(KEY4);
+    requestedKeys.add(KEY5);
+    Iterator<Entity> iter = createDatastoreForDeferredLookup().get(KEY1, KEY2, KEY3, KEY4, KEY5);
+    Set<Key> keysOfFoundEntities = new HashSet<>();
+    while (iter.hasNext()) {
+      keysOfFoundEntities.add(iter.next().key());
+    }
+    assertEquals(requestedKeys, keysOfFoundEntities);
+  }
+
+  public void testFetchArrayDeferredResults() throws DatastoreRpcException {
+    List<Entity> foundEntities =
+        createDatastoreForDeferredLookup().fetch(KEY1, KEY2, KEY3, KEY4, KEY5);
+    assertEquals(foundEntities.get(0).key(), KEY1);
+    assertEquals(foundEntities.get(1).key(), KEY2);
+    assertEquals(foundEntities.get(2).key(), KEY3);
+    assertEquals(foundEntities.get(3).key(), KEY4);
+    assertEquals(foundEntities.get(4).key(), KEY5);
+    assertEquals(foundEntities.size(), 5);
+  }
+
+  private Datastore createDatastoreForDeferredLookup() throws DatastoreRpcException {
+    List<DatastoreV1.Key> keysPb = new ArrayList<>();
+    keysPb.add(KEY1.toPb());
+    keysPb.add(KEY2.toPb());
+    keysPb.add(KEY3.toPb());
+    keysPb.add(KEY4.toPb());
+    keysPb.add(KEY5.toPb());
+    List<DatastoreV1.LookupRequest> lookupRequests = new ArrayList<>();
+    lookupRequests.add(DatastoreV1.LookupRequest.newBuilder().addAllKey(keysPb).build());
+    lookupRequests.add(
+        DatastoreV1.LookupRequest.newBuilder()
+            .addKey(keysPb.get(2))
+            .addKey(keysPb.get(3))
+            .addKey(keysPb.get(5))
+            .build());
+    lookupRequests.add(DatastoreV1.LookupRequest.newBuilder().addKey(keysPb.get(5)).build());
+    Entity entity4 = Entity.builder(KEY4).set("value", StringValue.of("value")).build();
+    Entity entity5 = Entity.builder(KEY5).set("value", "value").build();
+    List<DatastoreV1.LookupResponse> lookupResponses = new ArrayList<>();
+    lookupResponses.add(
+        DatastoreV1.LookupResponse.newBuilder()
+            .addFound(EntityResult.newBuilder().setEntity(ENTITY1.toPb()))
+            .addFound(EntityResult.newBuilder().setEntity(entity4.toPb()))
+            .addDeferred(keysPb.get(2))
+            .addDeferred(keysPb.get(3))
+            .addDeferred(keysPb.get(5))
+            .build());
+    lookupResponses.add(
+        DatastoreV1.LookupResponse.newBuilder()
+            .addFound(EntityResult.newBuilder().setEntity(ENTITY3.toPb()))
+            .addFound(EntityResult.newBuilder().setEntity(entity4.toPb()))
+            .addDeferred(keysPb.get(5))
+            .build());
+    lookupResponses.add(
+        DatastoreV1.LookupResponse.newBuilder()
+            .addFound(EntityResult.newBuilder().setEntity(entity5.toPb()))
+            .build());
+    DatastoreRpcFactory rpcFactoryMock = EasyMock.createStrictMock(DatastoreRpcFactory.class);
+    DatastoreRpc rpcMock = EasyMock.createStrictMock(DatastoreRpc.class);
+    EasyMock.expect(rpcFactoryMock.create(EasyMock.anyObject(DatastoreOptions.class)))
+        .andReturn(rpcMock);
+    for (int i = 0; i < lookupRequests.size(); i++) {
+      EasyMock.expect(rpcMock.lookup(lookupRequests.get(i))).andReturn(lookupResponses.get(i));
+    }
+    EasyMock.replay(rpcFactoryMock, rpcMock);
+    DatastoreOptions options =
+        this.options.toBuilder()
+            .retryParams(RetryParams.getDefaultInstance())
+            .serviceRpcFactory(rpcFactoryMock)
+            .build();
+    return DatastoreFactory.instance().get(options);
   }
 
   @Test