|
35 | 35 | import com.google.gcloud.datastore.StructuredQuery.PropertyFilter;
|
36 | 36 | import com.google.gcloud.spi.DatastoreRpc;
|
37 | 37 | import com.google.gcloud.spi.DatastoreRpc.DatastoreRpcException.Reason;
|
| 38 | +import com.google.gcloud.spi.DatastoreRpcFactory; |
38 | 39 |
|
39 | 40 | import org.easymock.EasyMock;
|
40 | 41 | import org.junit.AfterClass;
|
@@ -637,18 +638,21 @@ public void testRetires() throws Exception {
|
637 | 638 | DatastoreV1.LookupRequest.newBuilder().addKey(KEY1.toPb()).build();
|
638 | 639 | DatastoreV1.LookupResponse responsePb = DatastoreV1.LookupResponse.newBuilder()
|
639 | 640 | .addFound(EntityResult.newBuilder().setEntity(ENTITY1.toPb())).build();
|
640 |
| - DatastoreRpc mock = EasyMock.createStrictMock(DatastoreRpc.class); |
641 |
| - EasyMock.expect(mock.lookup(requestPb)) |
| 641 | + DatastoreRpcFactory rpcFactoryMock = EasyMock.createStrictMock(DatastoreRpcFactory.class); |
| 642 | + DatastoreRpc rpcMock = EasyMock.createStrictMock(DatastoreRpc.class); |
| 643 | + EasyMock.expect(rpcFactoryMock.create(EasyMock.anyObject(DatastoreServiceOptions.class))) |
| 644 | + .andReturn(rpcMock); |
| 645 | + EasyMock.expect(rpcMock.lookup(requestPb)) |
642 | 646 | .andThrow(new DatastoreRpc.DatastoreRpcException(Reason.UNAVAILABLE))
|
643 | 647 | .andReturn(responsePb);
|
644 |
| - EasyMock.replay(mock); |
| 648 | + EasyMock.replay(rpcFactoryMock, rpcMock); |
645 | 649 | DatastoreServiceOptions options = this.options.toBuilder()
|
646 | 650 | .retryParams(RetryParams.getDefaultInstance())
|
647 |
| - .datastoreRpc(mock) |
| 651 | + .serviceRpcFactory(rpcFactoryMock) |
648 | 652 | .build();
|
649 | 653 | DatastoreService datastore = DatastoreServiceFactory.instance().get(options);
|
650 | 654 | Entity entity = datastore.get(KEY1);
|
651 | 655 | assertEquals(ENTITY1, entity);
|
652 |
| - EasyMock.verify(mock); |
| 656 | + EasyMock.verify(rpcFactoryMock, rpcMock); |
653 | 657 | }
|
654 | 658 | }
|
0 commit comments