diff --git a/gcloud-java-datastore/pom.xml b/gcloud-java-datastore/pom.xml index b33f66a19682..ee633fd57883 100644 --- a/gcloud-java-datastore/pom.xml +++ b/gcloud-java-datastore/pom.xml @@ -42,18 +42,6 @@ datastore-v1beta3-proto-client 0.0.1-SNAPSHOT - - com.google.apis - google-api-services-datastore-protobuf - v1beta2-rev1-2.1.2 - compile - - - com.google.api-client - google-api-client - - - junit junit diff --git a/gcloud-java-datastore/src/main/java/com/google/gcloud/datastore/Datastore.java b/gcloud-java-datastore/src/main/java/com/google/gcloud/datastore/Datastore.java index b49db1cacdfe..c0efaa52b4e8 100644 --- a/gcloud-java-datastore/src/main/java/com/google/gcloud/datastore/Datastore.java +++ b/gcloud-java-datastore/src/main/java/com/google/gcloud/datastore/Datastore.java @@ -18,6 +18,7 @@ import com.google.gcloud.Service; +import java.util.Iterator; import java.util.List; /** @@ -32,7 +33,6 @@ public interface Datastore extends Service, DatastoreReaderWri */ Transaction newTransaction(); - /** * A callback for running with a transactional * {@link com.google.gcloud.datastore.DatastoreReaderWriter}. @@ -45,7 +45,6 @@ interface TransactionCallable { T run(DatastoreReaderWriter readerWriter) throws Exception; } - /** * Invokes the callback's {@link Datastore.TransactionCallable#run} method with a * {@link DatastoreReaderWriter} that is associated with a new transaction. @@ -105,4 +104,39 @@ interface TransactionCallable { * Returns a new KeyFactory for this service */ KeyFactory newKeyFactory(); + + /** + * Returns an {@link Entity} for the given {@link Key} or {@code null} if it doesn't exist. + * {@link ReadOption}s can be specified if desired. + * + * @throws DatastoreException upon failure + */ + Entity get(Key key, ReadOption... options); + + /** + * Returns an {@link Entity} for each given {@link Key} that exists in the Datastore. The order of + * the result is unspecified. Results are loaded lazily, so it is possible to get a + * {@code DatastoreException} from the returned {@code Iterator}'s + * {@link Iterator#hasNext hasNext} or {@link Iterator#next next} methods. {@link ReadOption}s can + * be specified if desired. + * + * @throws DatastoreException upon failure + * @see #get(Key) + */ + Iterator get(Iterable keys, ReadOption... options); + + /** + * Returns a list with a value for each given key (ordered by input). {@code null} values are + * returned for nonexistent keys. When possible prefer using {@link #get(Key...)} to avoid eagerly + * loading the results. {@link ReadOption}s can be specified if desired. + */ + List fetch(Iterable keys, ReadOption... options); + + /** + * Submits a {@link Query} and returns its result. {@link ReadOption}s can be specified if + * desired. + * + * @throws DatastoreException upon failure + */ + QueryResults run(Query query, ReadOption... options); } diff --git a/gcloud-java-datastore/src/main/java/com/google/gcloud/datastore/DatastoreHelper.java b/gcloud-java-datastore/src/main/java/com/google/gcloud/datastore/DatastoreHelper.java index 5ea0f755d2f8..e3cf9c055576 100644 --- a/gcloud-java-datastore/src/main/java/com/google/gcloud/datastore/DatastoreHelper.java +++ b/gcloud-java-datastore/src/main/java/com/google/gcloud/datastore/DatastoreHelper.java @@ -20,6 +20,8 @@ import com.google.common.collect.Maps; import java.util.ArrayList; +import java.util.Arrays; +import java.util.Collections; import java.util.Iterator; import java.util.List; import java.util.Map; @@ -33,13 +35,16 @@ class DatastoreHelper { private DatastoreHelper() { } - static Key allocateId(Datastore service, IncompleteKey key) { return service.allocateId(new IncompleteKey[]{key}).get(0); } - static Entity get(DatastoreReader reader, Key key) { - return Iterators.getNext(reader.get(new Key[]{key}), null); + static Entity get(Transaction reader, Key key) { + return Iterators.getNext(reader.get(new Key[] {key}), null); + } + + static Entity get(Datastore reader, Key key, ReadOption... options) { + return Iterators.getNext(reader.get(Collections.singletonList(key), options), null); } static Entity add(DatastoreWriter writer, FullEntity entity) { @@ -51,11 +56,22 @@ static KeyFactory newKeyFactory(DatastoreOptions options) { } /** - * Returns a list with a value for each given key (ordered by input). - * A {@code null} would be returned for non-existing keys. + * Returns a list with a value for each given key (ordered by input). {@code null} values are + * returned for nonexistent keys. */ - static List fetch(DatastoreReader reader, Key... keys) { - Iterator entities = reader.get(keys); + static List fetch(Transaction reader, Key... keys) { + return compileEntities(keys, reader.get(keys)); + } + + /** + * Returns a list with a value for each given key (ordered by input). {@code null} values are + * returned for nonexistent keys. + */ + static List fetch(Datastore reader, Key[] keys, ReadOption... options) { + return compileEntities(keys, reader.get(Arrays.asList(keys), options)); + } + + private static List compileEntities(Key[] keys, Iterator entities) { Map map = Maps.newHashMapWithExpectedSize(keys.length); while (entities.hasNext()) { Entity entity = entities.next(); @@ -63,7 +79,7 @@ static List fetch(DatastoreReader reader, Key... keys) { } List list = new ArrayList<>(keys.length); for (Key key : keys) { - // this will include nulls for non-existing keys + // this will include nulls for nonexistent keys list.add(map.get(key)); } return list; diff --git a/gcloud-java-datastore/src/main/java/com/google/gcloud/datastore/DatastoreImpl.java b/gcloud-java-datastore/src/main/java/com/google/gcloud/datastore/DatastoreImpl.java index df4e22149fd5..80cf59980de1 100644 --- a/gcloud-java-datastore/src/main/java/com/google/gcloud/datastore/DatastoreImpl.java +++ b/gcloud-java-datastore/src/main/java/com/google/gcloud/datastore/DatastoreImpl.java @@ -20,11 +20,14 @@ import com.google.common.base.Preconditions; import com.google.common.collect.AbstractIterator; import com.google.common.collect.ImmutableList; +import com.google.common.collect.Iterables; import com.google.common.collect.Sets; +import com.google.datastore.v1beta3.ReadOptions.ReadConsistency; import com.google.gcloud.BaseService; import com.google.gcloud.RetryHelper; import com.google.gcloud.RetryHelper.RetryHelperException; import com.google.gcloud.RetryParams; +import com.google.gcloud.datastore.ReadOption.EventualConsistency; import com.google.gcloud.spi.DatastoreRpc; import com.google.protobuf.ByteString; @@ -70,6 +73,11 @@ public QueryResults run(Query query) { return run(null, query); } + @Override + public QueryResults run(Query query, ReadOption... options) { + return run(toReadOptionsPb(options), query); + } + QueryResults run(com.google.datastore.v1beta3.ReadOptions readOptionsPb, Query query) { return new QueryResultsImpl<>(this, readOptionsPb, query); } @@ -185,16 +193,42 @@ public Entity get(Key key) { return DatastoreHelper.get(this, key); } + @Override + public Entity get(Key key, ReadOption... options) { + return DatastoreHelper.get(this, key, options); + } + @Override public Iterator get(Key... keys) { return get(null, keys); } + @Override + public Iterator get(Iterable keys, ReadOption... options) { + return get(toReadOptionsPb(options), Iterables.toArray(keys, Key.class)); + } + + private static com.google.datastore.v1beta3.ReadOptions toReadOptionsPb(ReadOption... options) { + com.google.datastore.v1beta3.ReadOptions readOptionsPb = null; + if (options != null + && ReadOption.asImmutableMap(options).containsKey(EventualConsistency.class)) { + readOptionsPb = com.google.datastore.v1beta3.ReadOptions.newBuilder() + .setReadConsistency(ReadConsistency.EVENTUAL) + .build(); + } + return readOptionsPb; + } + @Override public List fetch(Key... keys) { return DatastoreHelper.fetch(this, keys); } + @Override + public List fetch(Iterable keys, ReadOption... options) { + return DatastoreHelper.fetch(this, Iterables.toArray(keys, Key.class), options); + } + Iterator get(com.google.datastore.v1beta3.ReadOptions readOptionsPb, final Key... keys) { if (keys.length == 0) { return Collections.emptyIterator(); diff --git a/gcloud-java-datastore/src/main/java/com/google/gcloud/datastore/DatastoreReader.java b/gcloud-java-datastore/src/main/java/com/google/gcloud/datastore/DatastoreReader.java index 4852dd53e16c..3d6e5ec73243 100644 --- a/gcloud-java-datastore/src/main/java/com/google/gcloud/datastore/DatastoreReader.java +++ b/gcloud-java-datastore/src/main/java/com/google/gcloud/datastore/DatastoreReader.java @@ -25,18 +25,17 @@ public interface DatastoreReader { /** - * Returns an {@link Entity} for the given {@link Key} or {@code null} if does not exists. + * Returns an {@link Entity} for the given {@link Key} or {@code null} if it doesn't exist. * * @throws DatastoreException upon failure */ Entity get(Key key); /** - * Returns an {@link Entity} for each given {@link Key} that exists in the Datastore. - * The order of the result is unspecified. - * Results are loaded lazily therefore it is possible to get a {@code DatastoreException} - * from the returned {@code Iterator}'s {@link Iterator#hasNext hasNext} or - * {@link Iterator#next next} methods. + * Returns an {@link Entity} for each given {@link Key} that exists in the Datastore. The order of + * the result is unspecified. Results are loaded lazily, so it is possible to get a + * {@code DatastoreException} from the returned {@code Iterator}'s + * {@link Iterator#hasNext hasNext} or {@link Iterator#next next} methods. * * @throws DatastoreException upon failure * @see #get(Key) @@ -44,14 +43,14 @@ public interface DatastoreReader { Iterator get(Key... key); /** - * Returns a list with a value for each given key (ordered by input). - * A {@code null} would be returned for non-existing keys. - * When possible prefer using {@link #get(Key...)} which does not eagerly loads the results. + * Returns a list with a value for each given key (ordered by input). {@code null} values are + * returned for nonexistent keys. When possible prefer using {@link #get(Key...)} to avoid eagerly + * loading the results. */ List fetch(Key... keys); /** - * Submit a {@link Query} and returns its result. + * Submits a {@link Query} and returns its result. * * @throws DatastoreException upon failure */ diff --git a/gcloud-java-datastore/src/main/java/com/google/gcloud/datastore/ReadOption.java b/gcloud-java-datastore/src/main/java/com/google/gcloud/datastore/ReadOption.java new file mode 100644 index 000000000000..f0de06d1651d --- /dev/null +++ b/gcloud-java-datastore/src/main/java/com/google/gcloud/datastore/ReadOption.java @@ -0,0 +1,68 @@ +/* + * Copyright 2016 Google Inc. All Rights Reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package com.google.gcloud.datastore; + +import com.google.common.collect.ImmutableMap; + +import java.io.Serializable; +import java.util.Map; + +/** + * Specifies options for read operations in Datastore, namely getting/fetching entities and running + * queries. + */ +public abstract class ReadOption implements Serializable { + + private static final long serialVersionUID = -4406964829189800528L; + + /** + * Specifies eventual consistency for reads from Datastore. Lookups and ancestor queries using + * this option permit Datastore to return stale results. + */ + public static final class EventualConsistency extends ReadOption { + + private static final long serialVersionUID = -6959530217724666172L; + + private final boolean eventualConsistency; + + private EventualConsistency(boolean eventualConsistency) { + this.eventualConsistency = eventualConsistency; + } + + public boolean isEventual() { + return eventualConsistency; + } + } + + private ReadOption() {} + + /** + * Returns a {@code ReadOption} that specifies eventual consistency, allowing Datastore to return + * stale results from gets, fetches, and ancestor queries. + */ + public static EventualConsistency eventualConsistency() { + return new EventualConsistency(true); + } + + static Map, ReadOption> asImmutableMap(ReadOption... options) { + ImmutableMap.Builder, ReadOption> builder = ImmutableMap.builder(); + for (ReadOption option : options) { + builder.put(option.getClass(), option); + } + return builder.build(); + } +} diff --git a/gcloud-java-datastore/src/main/java/com/google/gcloud/datastore/TransactionImpl.java b/gcloud-java-datastore/src/main/java/com/google/gcloud/datastore/TransactionImpl.java index 5a422927005a..469c14e1c78a 100644 --- a/gcloud-java-datastore/src/main/java/com/google/gcloud/datastore/TransactionImpl.java +++ b/gcloud-java-datastore/src/main/java/com/google/gcloud/datastore/TransactionImpl.java @@ -40,7 +40,7 @@ static class ResponseImpl implements Transaction.Response { @Override public List generatedKeys() { - Iterator results = + Iterator results = response.getMutationResultsList().iterator(); List generated = new ArrayList<>(numAutoAllocatedIds); for (int i = 0; i < numAutoAllocatedIds; i++) { @@ -66,7 +66,7 @@ public Entity get(Key key) { @Override public Iterator get(Key... keys) { validateActive(); - com.google.datastore.v1beta3.ReadOptions.Builder readOptionsPb = + com.google.datastore.v1beta3.ReadOptions.Builder readOptionsPb = com.google.datastore.v1beta3.ReadOptions.newBuilder(); readOptionsPb.setTransaction(transaction); return datastore.get(readOptionsPb.build(), keys); @@ -81,7 +81,7 @@ public List fetch(Key... keys) { @Override public QueryResults run(Query query) { validateActive(); - com.google.datastore.v1beta3.ReadOptions.Builder readOptionsPb = + com.google.datastore.v1beta3.ReadOptions.Builder readOptionsPb = com.google.datastore.v1beta3.ReadOptions.newBuilder(); readOptionsPb.setTransaction(transaction); return datastore.run(readOptionsPb.build(), query); @@ -91,7 +91,7 @@ public QueryResults run(Query query) { public Transaction.Response commit() { validateActive(); List mutationsPb = toMutationPbList(); - com.google.datastore.v1beta3.CommitRequest.Builder requestPb = + com.google.datastore.v1beta3.CommitRequest.Builder requestPb = com.google.datastore.v1beta3.CommitRequest.newBuilder(); requestPb.setMode(com.google.datastore.v1beta3.CommitRequest.Mode.TRANSACTIONAL); requestPb.setTransaction(transaction); diff --git a/gcloud-java-datastore/src/test/java/com/google/gcloud/datastore/DatastoreHelperTest.java b/gcloud-java-datastore/src/test/java/com/google/gcloud/datastore/DatastoreHelperTest.java index 55c8d0cf3ce6..61b266a9abc2 100644 --- a/gcloud-java-datastore/src/test/java/com/google/gcloud/datastore/DatastoreHelperTest.java +++ b/gcloud-java-datastore/src/test/java/com/google/gcloud/datastore/DatastoreHelperTest.java @@ -27,6 +27,7 @@ import static org.junit.Assert.assertTrue; import static org.junit.Assert.fail; +import com.google.common.collect.ImmutableList; import com.google.common.collect.Iterators; import com.google.gcloud.datastore.Datastore.TransactionCallable; @@ -66,52 +67,100 @@ public void testAllocateId() throws Exception { } @Test - public void testGet() throws Exception { + public void testGetWithDatastore() throws Exception { Datastore datastore = createStrictMock(Datastore.class); IncompleteKey pKey1 = IncompleteKey.builder("ds", "k").build(); Key key1 = Key.builder(pKey1, 1).build(); Entity entity1 = Entity.builder(key1).build(); Key key2 = Key.builder(pKey1, 2).build(); - expect(datastore.get(new Key[]{key1})) + ReadOption eventualConsistency = ReadOption.eventualConsistency(); + expect(datastore.get(Collections.singletonList(key1))) .andReturn(Collections.singletonList(entity1).iterator()); - expect(datastore.get(new Key[]{key2})) + expect(datastore.get(Collections.singletonList(key2))) .andReturn(Collections.emptyIterator()); + expect(datastore.get(Collections.singletonList(key1), eventualConsistency)) + .andReturn(Collections.singletonList(entity1).iterator()); replay(datastore); assertEquals(entity1, DatastoreHelper.get(datastore, key1)); assertNull(DatastoreHelper.get(datastore, key2)); + assertEquals(entity1, DatastoreHelper.get(datastore, key1, eventualConsistency)); verify(datastore); } + @Test + public void testGetWithTransaction() throws Exception { + Transaction transaction = createStrictMock(Transaction.class); + IncompleteKey pKey1 = IncompleteKey.builder("ds", "k").build(); + Key key1 = Key.builder(pKey1, 1).build(); + Entity entity1 = Entity.builder(key1).build(); + Key key2 = Key.builder(pKey1, 2).build(); + expect(transaction.get(new Key[] {key1})) + .andReturn(Collections.singletonList(entity1).iterator()); + expect(transaction.get(new Key[] {key2})).andReturn(Collections.emptyIterator()); + replay(transaction); + assertEquals(entity1, DatastoreHelper.get(transaction, key1)); + assertNull(DatastoreHelper.get(transaction, key2)); + verify(transaction); + } + @Test public void testAdd() throws Exception { Datastore datastore = createStrictMock(Datastore.class); IncompleteKey pKey = IncompleteKey.builder("ds", "k").build(); Key key = Key.builder(pKey, 1).build(); Entity entity = Entity.builder(key).build(); - expect(datastore.add(new Entity[]{entity})) - .andReturn(Collections.singletonList(entity)); + expect(datastore.add(new Entity[] {entity})).andReturn(Collections.singletonList(entity)); replay(datastore); assertEquals(entity, DatastoreHelper.add(datastore, entity)); verify(datastore); } @Test - public void testFetch() throws Exception { + public void testFetchWithDatastore() throws Exception { Datastore datastore = createStrictMock(Datastore.class); IncompleteKey pKey1 = IncompleteKey.builder("ds", "k").build(); Key key1 = Key.builder(pKey1, 1).build(); Key key2 = Key.builder(pKey1, "a").build(); Entity entity1 = Entity.builder(key1).build(); Entity entity2 = Entity.builder(key2).build(); - expect(datastore.get(key1, key2)).andReturn(Iterators.forArray(entity1, entity2)).once(); + ReadOption eventualConsistency = ReadOption.eventualConsistency(); + expect(datastore.get(ImmutableList.of(key1, key2))) + .andReturn(Iterators.forArray(entity1, entity2)) + .once(); + expect(datastore.get(ImmutableList.of(key1, key2), eventualConsistency)) + .andReturn(Iterators.forArray(entity1, entity2)) + .once(); replay(datastore); - List values = DatastoreHelper.fetch(datastore, key1, key2); + List values = DatastoreHelper.fetch(datastore, new Key[] {key1, key2}); + assertEquals(2, values.size()); + assertEquals(entity1, values.get(0)); + assertEquals(entity2, values.get(1)); + values = DatastoreHelper.fetch(datastore, new Key[] {key1, key2}, eventualConsistency); assertEquals(2, values.size()); assertEquals(entity1, values.get(0)); assertEquals(entity2, values.get(1)); verify(datastore); } + @Test + public void testFetchWithTransaction() throws Exception { + Transaction transaction = createStrictMock(Transaction.class); + IncompleteKey pKey1 = IncompleteKey.builder("ds", "k").build(); + Key key1 = Key.builder(pKey1, 1).build(); + Key key2 = Key.builder(pKey1, "a").build(); + Entity entity1 = Entity.builder(key1).build(); + Entity entity2 = Entity.builder(key2).build(); + expect(transaction.get(new Key[] {key1, key2})) + .andReturn(Iterators.forArray(entity1, entity2)) + .once(); + replay(transaction); + List values = DatastoreHelper.fetch(transaction, new Key[] {key1, key2}); + assertEquals(2, values.size()); + assertEquals(entity1, values.get(0)); + assertEquals(entity2, values.get(1)); + verify(transaction); + } + @Test public void testRunInTransaction() throws Exception { final Datastore datastore = createStrictMock(Datastore.class); 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 002e5f6df04f..23c22ac4f645 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 @@ -25,7 +25,17 @@ import static org.junit.Assert.assertTrue; import static org.junit.Assert.fail; +import com.google.common.collect.ImmutableList; import com.google.common.collect.Iterators; +import com.google.datastore.v1beta3.EntityResult; +import com.google.datastore.v1beta3.LookupRequest; +import com.google.datastore.v1beta3.LookupResponse; +import com.google.datastore.v1beta3.PartitionId; +import com.google.datastore.v1beta3.QueryResultBatch; +import com.google.datastore.v1beta3.ReadOptions; +import com.google.datastore.v1beta3.ReadOptions.ReadConsistency; +import com.google.datastore.v1beta3.RunQueryRequest; +import com.google.datastore.v1beta3.RunQueryResponse; import com.google.gcloud.RetryParams; import com.google.gcloud.datastore.Query.ResultType; import com.google.gcloud.datastore.StructuredQuery.OrderBy; @@ -91,22 +101,24 @@ public class DatastoreTest { private static final FullEntity PARTIAL_ENTITY3 = FullEntity.builder(PARTIAL_ENTITY1).key(IncompleteKey.builder(PROJECT_ID, KIND3).build()) .build(); - private static final Entity ENTITY1 = - Entity.builder(KEY1) - .set("str", STR_VALUE) - .set("date", DATE_TIME_VALUE) - .set("latLng", LAT_LNG_VALUE) - .set("bool", BOOL_VALUE) - .set("partial1", EntityValue.of(PARTIAL_ENTITY1)) - .set("list", LIST_VALUE2) - .build(); + private static final Entity ENTITY1 = Entity.builder(KEY1) + .set("str", STR_VALUE) + .set("date", DATE_TIME_VALUE) + .set("latLng", LAT_LNG_VALUE) + .set("bool", BOOL_VALUE) + .set("partial1", EntityValue.of(PARTIAL_ENTITY1)) + .set("list", LIST_VALUE2) + .build(); private static final Entity ENTITY2 = Entity.builder(ENTITY1).key(KEY2).remove("str") .set("name", "Dan").setNull("null").set("age", 20).build(); private static final Entity ENTITY3 = Entity.builder(ENTITY1).key(KEY3).remove("str") .set("null", NULL_VALUE).set("partial1", PARTIAL_ENTITY2).set("partial2", ENTITY2).build(); private DatastoreOptions options; + private DatastoreOptions rpcMockOptions; private Datastore datastore; + private DatastoreRpcFactory rpcFactoryMock; + private DatastoreRpc rpcMock; private static LocalGcdHelper gcdHelper; private static final int PORT = LocalGcdHelper.findAvailablePort(LocalGcdHelper.DEFAULT_PORT); @@ -129,6 +141,14 @@ public void setUp() { .retryParams(RetryParams.noRetries()) .build(); datastore = options.service(); + rpcFactoryMock = EasyMock.createStrictMock(DatastoreRpcFactory.class); + rpcMock = EasyMock.createStrictMock(DatastoreRpc.class); + rpcMockOptions = options + .toBuilder() + .retryParams(RetryParams.defaultInstance()) + .serviceRpcFactory(rpcFactoryMock) + .build(); + EasyMock.expect(rpcFactoryMock.create(rpcMockOptions)).andReturn(rpcMock); StructuredQuery query = Query.keyQueryBuilder().build(); QueryResults result = datastore.run(query); datastore.delete(Iterators.toArray(result, Key.class)); @@ -421,24 +441,13 @@ public void testRunGqlQueryWithCasting() { @Test public void testGqlQueryPagination() throws DatastoreException { - DatastoreRpcFactory rpcFactoryMock = EasyMock.createStrictMock(DatastoreRpcFactory.class); - DatastoreRpc rpcMock = EasyMock.createStrictMock(DatastoreRpc.class); - EasyMock.expect(rpcFactoryMock.create(EasyMock.anyObject(DatastoreOptions.class))) - .andReturn(rpcMock); - List responses = - buildResponsesForQueryPagination(); + List responses = buildResponsesForQueryPagination(); for (int i = 0; i < responses.size(); i++) { - EasyMock - .expect(rpcMock.runQuery( - EasyMock.anyObject(com.google.datastore.v1beta3.RunQueryRequest.class))) - .andReturn(responses.get(i)); + EasyMock.expect(rpcMock.runQuery(EasyMock.anyObject(RunQueryRequest.class))) + .andReturn(responses.get(i)); } EasyMock.replay(rpcFactoryMock, rpcMock); - DatastoreOptions options = this.options.toBuilder() - .retryParams(RetryParams.defaultInstance()) - .serviceRpcFactory(rpcFactoryMock) - .build(); - Datastore mockDatastore = options.service(); + Datastore mockDatastore = rpcMockOptions.service(); QueryResults results = mockDatastore.run(Query.gqlQueryBuilder(ResultType.KEY, "select __key__ from *").build()); int count = 0; @@ -475,15 +484,14 @@ public void testRunStructuredQuery() { assertTrue(projectionEntity.names().isEmpty()); assertFalse(results2.hasNext()); - StructuredQuery projectionQuery = - Query.projectionEntityQueryBuilder() - .kind(KIND2) - .projection("age") - .filter(PropertyFilter.gt("age", 18)) - .distinctOn("age") - .orderBy(OrderBy.asc("age")) - .limit(10) - .build(); + StructuredQuery projectionQuery = Query.projectionEntityQueryBuilder() + .kind(KIND2) + .projection("age") + .filter(PropertyFilter.gt("age", 18)) + .distinctOn("age") + .orderBy(OrderBy.asc("age")) + .limit(10) + .build(); QueryResults results4 = datastore.run(projectionQuery); assertTrue(results4.hasNext()); @@ -496,25 +504,14 @@ public void testRunStructuredQuery() { @Test public void testStructuredQueryPagination() throws DatastoreException { - DatastoreRpcFactory rpcFactoryMock = EasyMock.createStrictMock(DatastoreRpcFactory.class); - DatastoreRpc rpcMock = EasyMock.createStrictMock(DatastoreRpc.class); - EasyMock.expect(rpcFactoryMock.create(EasyMock.anyObject(DatastoreOptions.class))) - .andReturn(rpcMock); - List responses = - buildResponsesForQueryPagination(); + List responses = buildResponsesForQueryPagination(); for (int i = 0; i < responses.size(); i++) { - EasyMock - .expect(rpcMock.runQuery( - EasyMock.anyObject(com.google.datastore.v1beta3.RunQueryRequest.class))) - .andReturn(responses.get(i)); + EasyMock.expect(rpcMock.runQuery(EasyMock.anyObject(RunQueryRequest.class))) + .andReturn(responses.get(i)); } EasyMock.replay(rpcFactoryMock, rpcMock); - DatastoreOptions options = this.options.toBuilder() - .retryParams(RetryParams.defaultInstance()) - .serviceRpcFactory(rpcFactoryMock) - .build(); - Datastore mockDatastore = options.service(); - QueryResults results = mockDatastore.run(Query.keyQueryBuilder().build()); + Datastore datastore = rpcMockOptions.service(); + QueryResults results = datastore.run(Query.keyQueryBuilder().build()); int count = 0; while (results.hasNext()) { count += 1; @@ -524,85 +521,57 @@ public void testStructuredQueryPagination() throws DatastoreException { EasyMock.verify(rpcFactoryMock, rpcMock); } - private List buildResponsesForQueryPagination() { + private List 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 responses = new ArrayList<>(); + List responses = new ArrayList<>(); Query query = Query.keyQueryBuilder().build(); - com.google.datastore.v1beta3.RunQueryRequest.Builder requestPb = - com.google.datastore.v1beta3.RunQueryRequest.newBuilder(); + RunQueryRequest.Builder requestPb = RunQueryRequest.newBuilder(); query.populatePb(requestPb); - com.google.datastore.v1beta3.QueryResultBatch queryResultBatchPb = - com.google.datastore.v1beta3.RunQueryResponse.newBuilder() - .mergeFrom(((DatastoreImpl) datastore).runQuery(requestPb.build())) - .getBatch(); - com.google.datastore.v1beta3.QueryResultBatch queryResultBatchPb1 = - com.google.datastore.v1beta3.QueryResultBatch.newBuilder() - .mergeFrom(queryResultBatchPb) - .setMoreResults( - com.google.datastore.v1beta3.QueryResultBatch.MoreResultsType.NOT_FINISHED) - .clearEntityResults() - .addAllEntityResults(queryResultBatchPb.getEntityResultsList().subList(0, 1)) - .setEndCursor(queryResultBatchPb.getEntityResultsList().get(0).getCursor()) - .build(); - responses.add( - com.google.datastore.v1beta3.RunQueryResponse.newBuilder() - .setBatch(queryResultBatchPb1) - .build()); - com.google.datastore.v1beta3.QueryResultBatch queryResultBatchPb2 = - com.google.datastore.v1beta3.QueryResultBatch.newBuilder() - .mergeFrom(queryResultBatchPb) - .setMoreResults( - com.google.datastore.v1beta3.QueryResultBatch.MoreResultsType.NOT_FINISHED) - .clearEntityResults() - .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()); + QueryResultBatch queryResultBatchPb = RunQueryResponse.newBuilder() + .mergeFrom(((DatastoreImpl) datastore).runQuery(requestPb.build())) + .getBatch(); + QueryResultBatch queryResultBatchPb1 = QueryResultBatch.newBuilder() + .mergeFrom(queryResultBatchPb) + .setMoreResults(QueryResultBatch.MoreResultsType.NOT_FINISHED) + .clearEntityResults() + .addAllEntityResults(queryResultBatchPb.getEntityResultsList().subList(0, 1)) + .setEndCursor(queryResultBatchPb.getEntityResultsList().get(0).getCursor()) + .build(); + responses.add(RunQueryResponse.newBuilder().setBatch(queryResultBatchPb1).build()); + QueryResultBatch queryResultBatchPb2 = QueryResultBatch.newBuilder() + .mergeFrom(queryResultBatchPb) + .setMoreResults(QueryResultBatch.MoreResultsType.NOT_FINISHED) + .clearEntityResults() + .addAllEntityResults(queryResultBatchPb.getEntityResultsList().subList(1, 3)) + .setEndCursor(queryResultBatchPb.getEntityResultsList().get(2).getCursor()) + .build(); + responses.add(RunQueryResponse.newBuilder().setBatch(queryResultBatchPb2).build()); + QueryResultBatch queryResultBatchPb3 = QueryResultBatch.newBuilder() + .mergeFrom(queryResultBatchPb) + .setMoreResults(QueryResultBatch.MoreResultsType.NO_MORE_RESULTS) + .clearEntityResults() + .addAllEntityResults(queryResultBatchPb.getEntityResultsList().subList(3, 5)) + .setEndCursor(queryResultBatchPb.getEntityResultsList().get(4).getCursor()) + .build(); + responses.add(RunQueryResponse.newBuilder().setBatch(queryResultBatchPb3).build()); return responses; } public void testQueryPaginationWithLimit() throws DatastoreException { - DatastoreRpcFactory rpcFactoryMock = EasyMock.createStrictMock(DatastoreRpcFactory.class); - DatastoreRpc rpcMock = EasyMock.createStrictMock(DatastoreRpc.class); - EasyMock.expect(rpcFactoryMock.create(EasyMock.anyObject(DatastoreOptions.class))) - .andReturn(rpcMock); - List responses = - buildResponsesForQueryPaginationWithLimit(); + List responses = buildResponsesForQueryPaginationWithLimit(); for (int i = 0; i < responses.size(); i++) { - EasyMock.expect( - rpcMock.runQuery( - EasyMock.anyObject(com.google.datastore.v1beta3.RunQueryRequest.class))) + EasyMock.expect(rpcMock.runQuery(EasyMock.anyObject(RunQueryRequest.class))) .andReturn(responses.get(i)); } EasyMock.replay(rpcFactoryMock, rpcMock); - Datastore mockDatastore = options.toBuilder() - .retryParams(RetryParams.defaultInstance()) - .serviceRpcFactory(rpcFactoryMock) - .build() - .service(); + Datastore datastore = rpcMockOptions.service(); int limit = 2; int totalCount = 0; StructuredQuery query = Query.entityQueryBuilder().limit(limit).build(); while (true) { - QueryResults results = mockDatastore.run(query); + QueryResults results = datastore.run(query); int resultCount = 0; while (results.hasNext()) { results.next(); @@ -618,78 +587,74 @@ public void testQueryPaginationWithLimit() throws DatastoreException { EasyMock.verify(rpcFactoryMock, rpcMock); } - private List - buildResponsesForQueryPaginationWithLimit() { + private List 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); - List responses = new ArrayList<>(); + List responses = new ArrayList<>(); Query query = Query.entityQueryBuilder().build(); - com.google.datastore.v1beta3.RunQueryRequest.Builder requestPb = - com.google.datastore.v1beta3.RunQueryRequest.newBuilder(); + RunQueryRequest.Builder requestPb = RunQueryRequest.newBuilder(); query.populatePb(requestPb); - com.google.datastore.v1beta3.QueryResultBatch queryResultBatchPb = - com.google.datastore.v1beta3.RunQueryResponse.newBuilder() - .mergeFrom(((DatastoreImpl) datastore).runQuery(requestPb.build())) - .getBatch(); - com.google.datastore.v1beta3.QueryResultBatch queryResultBatchPb1 = - com.google.datastore.v1beta3.QueryResultBatch.newBuilder() - .mergeFrom(queryResultBatchPb) - .setMoreResults( - com.google.datastore.v1beta3.QueryResultBatch.MoreResultsType.NOT_FINISHED) - .clearEntityResults() - .addAllEntityResults(queryResultBatchPb.getEntityResultsList().subList(0, 1)) - .setEndCursor(queryResultBatchPb.getEntityResultsList().get(0).getCursor()) - .build(); - responses.add( - com.google.datastore.v1beta3.RunQueryResponse.newBuilder() - .setBatch(queryResultBatchPb1) - .build()); - com.google.datastore.v1beta3.QueryResultBatch queryResultBatchPb2 = - com.google.datastore.v1beta3.QueryResultBatch.newBuilder() - .mergeFrom(queryResultBatchPb) - .setMoreResults( - com.google.datastore.v1beta3.QueryResultBatch.MoreResultsType - .MORE_RESULTS_AFTER_LIMIT) - .clearEntityResults() - .addAllEntityResults(queryResultBatchPb.getEntityResultsList().subList(1, 2)) - .setEndCursor( - ByteString.copyFrom(new byte[] {(byte) 0x80})) // test invalid UTF-8 string - .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 - .MORE_RESULTS_AFTER_LIMIT) - .clearEntityResults() - .addAllEntityResults(queryResultBatchPb.getEntityResultsList().subList(2, 4)) - .setEndCursor(queryResultBatchPb.getEntityResultsList().get(3).getCursor()) - .build(); - responses.add( - com.google.datastore.v1beta3.RunQueryResponse.newBuilder() - .setBatch(queryResultBatchPb3) - .build()); - com.google.datastore.v1beta3.QueryResultBatch queryResultBatchPb4 = - com.google.datastore.v1beta3.QueryResultBatch.newBuilder() - .mergeFrom(queryResultBatchPb) - .setMoreResults( - com.google.datastore.v1beta3.QueryResultBatch.MoreResultsType.NO_MORE_RESULTS) - .clearEntityResults() - .addAllEntityResults(queryResultBatchPb.getEntityResultsList().subList(4, 5)) - .setEndCursor(queryResultBatchPb.getEntityResultsList().get(4).getCursor()) - .build(); - responses.add( - com.google.datastore.v1beta3.RunQueryResponse.newBuilder() - .setBatch(queryResultBatchPb4) - .build()); + QueryResultBatch queryResultBatchPb = RunQueryResponse.newBuilder() + .mergeFrom(((DatastoreImpl) datastore).runQuery(requestPb.build())) + .getBatch(); + QueryResultBatch queryResultBatchPb1 = QueryResultBatch.newBuilder() + .mergeFrom(queryResultBatchPb) + .setMoreResults(QueryResultBatch.MoreResultsType.NOT_FINISHED) + .clearEntityResults() + .addAllEntityResults(queryResultBatchPb.getEntityResultsList().subList(0, 1)) + .setEndCursor(queryResultBatchPb.getEntityResultsList().get(0).getCursor()) + .build(); + responses.add(RunQueryResponse.newBuilder().setBatch(queryResultBatchPb1).build()); + QueryResultBatch queryResultBatchPb2 = QueryResultBatch.newBuilder() + .mergeFrom(queryResultBatchPb) + .setMoreResults(QueryResultBatch.MoreResultsType.MORE_RESULTS_AFTER_LIMIT) + .clearEntityResults() + .addAllEntityResults(queryResultBatchPb.getEntityResultsList().subList(1, 2)) + .setEndCursor( + ByteString.copyFrom(new byte[] {(byte) 0x80})) // test invalid UTF-8 string + .build(); + responses.add(RunQueryResponse.newBuilder().setBatch(queryResultBatchPb2).build()); + QueryResultBatch queryResultBatchPb3 = QueryResultBatch.newBuilder() + .mergeFrom(queryResultBatchPb) + .setMoreResults(QueryResultBatch.MoreResultsType.MORE_RESULTS_AFTER_LIMIT) + .clearEntityResults() + .addAllEntityResults(queryResultBatchPb.getEntityResultsList().subList(2, 4)) + .setEndCursor(queryResultBatchPb.getEntityResultsList().get(3).getCursor()) + .build(); + responses.add(RunQueryResponse.newBuilder().setBatch(queryResultBatchPb3).build()); + QueryResultBatch queryResultBatchPb4 = QueryResultBatch.newBuilder() + .mergeFrom(queryResultBatchPb) + .setMoreResults(QueryResultBatch.MoreResultsType.NO_MORE_RESULTS) + .clearEntityResults() + .addAllEntityResults(queryResultBatchPb.getEntityResultsList().subList(4, 5)) + .setEndCursor(queryResultBatchPb.getEntityResultsList().get(4).getCursor()) + .build(); + responses.add(RunQueryResponse.newBuilder().setBatch(queryResultBatchPb4).build()); return responses; } + @Test + public void testEventualConsistencyQuery() { + ReadOptions readOption = + ReadOptions.newBuilder().setReadConsistencyValue(ReadConsistency.EVENTUAL_VALUE).build(); + com.google.datastore.v1beta3.GqlQuery query = com.google.datastore.v1beta3.GqlQuery.newBuilder() + .setQueryString("FROM * SELECT *") + .build(); + RunQueryRequest.Builder expectedRequest = RunQueryRequest.newBuilder() + .setReadOptions(readOption) + .setGqlQuery(query) + .setPartitionId(PartitionId.newBuilder().setProjectId(PROJECT_ID) + .build()); + EasyMock.expect(rpcMock.runQuery(expectedRequest.build())) + .andReturn(RunQueryResponse.newBuilder().build()); + EasyMock.replay(rpcFactoryMock, rpcMock); + Datastore datastore = rpcMockOptions.service(); + datastore.run( + Query.gqlQueryBuilder("FROM * SELECT *").build(), ReadOption.eventualConsistency()); + EasyMock.verify(rpcFactoryMock, rpcMock); + } + @Test public void testToUrlSafe() { byte[][] invalidUtf8 = @@ -765,6 +730,28 @@ public void testGet() { assertFalse(entity.contains("bla")); } + @Test + public void testLookupEventualConsistency() { + ReadOptions readOption = + ReadOptions.newBuilder().setReadConsistencyValue(ReadConsistency.EVENTUAL_VALUE).build(); + com.google.datastore.v1beta3.Key key = com.google.datastore.v1beta3.Key.newBuilder() + .setPartitionId(PartitionId.newBuilder().setProjectId(PROJECT_ID).build()) + .addPath(com.google.datastore.v1beta3.Key.PathElement.newBuilder() + .setKind("kind1").setName("name").build()) + .build(); + LookupRequest lookupRequest = + LookupRequest.newBuilder().setReadOptions(readOption).addKeys(key).build(); + EasyMock.expect(rpcMock.lookup(lookupRequest)) + .andReturn(LookupResponse.newBuilder().build()) + .times(3); + EasyMock.replay(rpcFactoryMock, rpcMock); + Datastore datastore = rpcMockOptions.service(); + datastore.get(KEY1, ReadOption.eventualConsistency()); + datastore.get(ImmutableList.of(KEY1), ReadOption.eventualConsistency()); + datastore.fetch(ImmutableList.of(KEY1), ReadOption.eventualConsistency()); + EasyMock.verify(rpcFactoryMock, rpcMock); + } + @Test public void testGetArrayNoDeferredResults() { datastore.put(ENTITY3); @@ -828,57 +815,41 @@ private Datastore createDatastoreForDeferredLookup() throws DatastoreException { keysPb.add(KEY3.toPb()); keysPb.add(KEY4.toPb()); keysPb.add(KEY5.toPb()); - List lookupRequests = new ArrayList<>(); + List lookupRequests = new ArrayList<>(); + lookupRequests.add(LookupRequest.newBuilder().addAllKeys(keysPb).build()); lookupRequests.add( - com.google.datastore.v1beta3.LookupRequest.newBuilder().addAllKeys(keysPb).build()); - lookupRequests.add( - com.google.datastore.v1beta3.LookupRequest.newBuilder() + LookupRequest.newBuilder() .addKeys(keysPb.get(2)) .addKeys(keysPb.get(3)) .addKeys(keysPb.get(5)) .build()); - lookupRequests.add( - com.google.datastore.v1beta3.LookupRequest.newBuilder().addKeys(keysPb.get(5)).build()); + lookupRequests.add(LookupRequest.newBuilder().addKeys(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<>(); + List lookupResponses = new ArrayList<>(); lookupResponses.add( - com.google.datastore.v1beta3.LookupResponse.newBuilder() - .addFound( - com.google.datastore.v1beta3.EntityResult.newBuilder().setEntity(ENTITY1.toPb())) - .addFound( - com.google.datastore.v1beta3.EntityResult.newBuilder().setEntity(entity4.toPb())) + 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( - com.google.datastore.v1beta3.LookupResponse.newBuilder() - .addFound( - com.google.datastore.v1beta3.EntityResult.newBuilder().setEntity(ENTITY3.toPb())) - .addFound( - com.google.datastore.v1beta3.EntityResult.newBuilder().setEntity(entity4.toPb())) + LookupResponse.newBuilder() + .addFound(EntityResult.newBuilder().setEntity(ENTITY3.toPb())) + .addFound(EntityResult.newBuilder().setEntity(entity4.toPb())) .addDeferred(keysPb.get(5)) .build()); lookupResponses.add( - com.google.datastore.v1beta3.LookupResponse.newBuilder() - .addFound( - com.google.datastore.v1beta3.EntityResult.newBuilder().setEntity(entity5.toPb())) + 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.defaultInstance()) - .serviceRpcFactory(rpcFactoryMock) - .build(); - return options.service(); + return rpcMockOptions.service(); } @Test @@ -974,26 +945,15 @@ public void testKeyFactory() { @Test public void testRetryableException() throws Exception { - com.google.datastore.v1beta3.LookupRequest requestPb = - com.google.datastore.v1beta3.LookupRequest.newBuilder().addKeys(KEY1.toPb()).build(); - com.google.datastore.v1beta3.LookupResponse responsePb = - com.google.datastore.v1beta3.LookupResponse.newBuilder() - .addFound( - com.google.datastore.v1beta3.EntityResult.newBuilder().setEntity(ENTITY1.toPb())) - .build(); - DatastoreRpcFactory rpcFactoryMock = EasyMock.createStrictMock(DatastoreRpcFactory.class); - DatastoreRpc rpcMock = EasyMock.createStrictMock(DatastoreRpc.class); - EasyMock.expect(rpcFactoryMock.create(EasyMock.anyObject(DatastoreOptions.class))) - .andReturn(rpcMock); + LookupRequest requestPb = LookupRequest.newBuilder().addKeys(KEY1.toPb()).build(); + LookupResponse responsePb = LookupResponse.newBuilder() + .addFound(EntityResult.newBuilder().setEntity(ENTITY1.toPb())) + .build(); EasyMock.expect(rpcMock.lookup(requestPb)) .andThrow(new DatastoreException(14, "UNAVAILABLE", "UNAVAILABLE", null)) .andReturn(responsePb); EasyMock.replay(rpcFactoryMock, rpcMock); - DatastoreOptions options = this.options.toBuilder() - .retryParams(RetryParams.defaultInstance()) - .serviceRpcFactory(rpcFactoryMock) - .build(); - Datastore datastore = options.service(); + Datastore datastore = rpcMockOptions.service(); Entity entity = datastore.get(KEY1); assertEquals(ENTITY1, entity); EasyMock.verify(rpcFactoryMock, rpcMock); @@ -1001,23 +961,13 @@ public void testRetryableException() throws Exception { @Test public void testNonRetryableException() throws Exception { - com.google.datastore.v1beta3.LookupRequest requestPb = - com.google.datastore.v1beta3.LookupRequest.newBuilder().addKeys(KEY1.toPb()).build(); - DatastoreRpcFactory rpcFactoryMock = EasyMock.createStrictMock(DatastoreRpcFactory.class); - DatastoreRpc rpcMock = EasyMock.createStrictMock(DatastoreRpc.class); - EasyMock.expect(rpcFactoryMock.create(EasyMock.anyObject(DatastoreOptions.class))) - .andReturn(rpcMock); + LookupRequest requestPb = LookupRequest.newBuilder().addKeys(KEY1.toPb()).build(); EasyMock.expect(rpcMock.lookup(requestPb)) .andThrow( new DatastoreException(DatastoreException.UNKNOWN_CODE, "denied", "PERMISSION_DENIED")) .times(1); EasyMock.replay(rpcFactoryMock, rpcMock); - RetryParams retryParams = RetryParams.builder().retryMinAttempts(2).build(); - DatastoreOptions options = this.options.toBuilder() - .retryParams(retryParams) - .serviceRpcFactory(rpcFactoryMock) - .build(); - Datastore datastore = options.service(); + Datastore datastore = rpcMockOptions.service(); thrown.expect(DatastoreException.class); thrown.expectMessage("denied"); datastore.get(KEY1); @@ -1026,21 +976,11 @@ public void testNonRetryableException() throws Exception { @Test public void testRuntimeException() throws Exception { - com.google.datastore.v1beta3.LookupRequest requestPb = - com.google.datastore.v1beta3.LookupRequest.newBuilder().addKeys(KEY1.toPb()).build(); - DatastoreRpcFactory rpcFactoryMock = EasyMock.createStrictMock(DatastoreRpcFactory.class); - DatastoreRpc rpcMock = EasyMock.createStrictMock(DatastoreRpc.class); - EasyMock.expect(rpcFactoryMock.create(EasyMock.anyObject(DatastoreOptions.class))) - .andReturn(rpcMock); + LookupRequest requestPb = LookupRequest.newBuilder().addKeys(KEY1.toPb()).build(); String exceptionMessage = "Artificial runtime exception"; - EasyMock.expect(rpcMock.lookup(requestPb)) - .andThrow(new RuntimeException(exceptionMessage)); + EasyMock.expect(rpcMock.lookup(requestPb)).andThrow(new RuntimeException(exceptionMessage)); EasyMock.replay(rpcFactoryMock, rpcMock); - DatastoreOptions options = this.options.toBuilder() - .retryParams(RetryParams.defaultInstance()) - .serviceRpcFactory(rpcFactoryMock) - .build(); - Datastore datastore = options.service(); + Datastore datastore = rpcMockOptions.service(); thrown.expect(DatastoreException.class); thrown.expectMessage(exceptionMessage); datastore.get(KEY1);