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 0bef237ee48a..c55820b324fb 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 @@ -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 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 requestedKeys = new HashSet<>(); + requestedKeys.add(KEY1); + requestedKeys.add(KEY2); + requestedKeys.add(KEY3); + requestedKeys.add(KEY4); + requestedKeys.add(KEY5); + Iterator iter = createDatastoreForDeferredLookup().get(KEY1, KEY2, KEY3, KEY4, KEY5); + Set keysOfFoundEntities = new HashSet<>(); + while (iter.hasNext()) { + keysOfFoundEntities.add(iter.next().key()); + } + assertEquals(requestedKeys, keysOfFoundEntities); + } + + public void testFetchArrayDeferredResults() throws DatastoreRpcException { + List 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 keysPb = new ArrayList<>(); + keysPb.add(KEY1.toPb()); + keysPb.add(KEY2.toPb()); + keysPb.add(KEY3.toPb()); + keysPb.add(KEY4.toPb()); + keysPb.add(KEY5.toPb()); + List 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 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