@Test(
dataProvider =
com.linkedin.restli.internal.common.TestConstants.RESTLI_PROTOCOL_1_2_PREFIX
+ "requestOptionsDataProvider")
public void testAssociationBatchGetKVCompoundKeyResponse(RestliRequestOptions requestOptions)
throws RemoteInvocationException {
CompoundKey key1 = buildCompoundKey(1, 1);
CompoundKey key2 = buildCompoundKey(2, 1);
Set<CompoundKey> allRequestedKeys = new HashSet<CompoundKey>(Arrays.asList(key1, key2));
Request<BatchKVResponse<CompoundKey, GroupMembership>> request =
new GroupMembershipsBuilders(requestOptions)
.batchGet()
.ids(key1, key2)
.fields(GroupMembership.fields().contactEmail())
.buildKV();
BatchKVResponse<CompoundKey, GroupMembership> groupMemberships =
getClient().sendRequest(request).getResponse().getEntity();
Assert.assertTrue(allRequestedKeys.containsAll(groupMemberships.getResults().keySet()));
Assert.assertTrue(allRequestedKeys.containsAll(groupMemberships.getErrors().keySet()));
Set<CompoundKey> allResponseKeys =
new HashSet<CompoundKey>(groupMemberships.getResults().keySet());
allResponseKeys.addAll(groupMemberships.getErrors().keySet());
Assert.assertEquals(allResponseKeys, allRequestedKeys);
}
private static void testSendSGKVRequests(
Collection<ScatterGatherBuilder.KVRequestInfo<Long, UpdateStatus>> scatterGatherRequests,
Long[] requestIds)
throws InterruptedException {
final Map<Long, UpdateStatus> results = new ConcurrentHashMap<Long, UpdateStatus>();
final CountDownLatch latch = new CountDownLatch(scatterGatherRequests.size());
final List<Throwable> errors = new ArrayList<Throwable>();
final List<BatchKVResponse<Long, UpdateStatus>> responses =
new ArrayList<BatchKVResponse<Long, UpdateStatus>>();
for (ScatterGatherBuilder.KVRequestInfo<Long, UpdateStatus> requestInfo :
scatterGatherRequests) {
Callback<Response<BatchKVResponse<Long, UpdateStatus>>> cb =
new Callback<Response<BatchKVResponse<Long, UpdateStatus>>>() {
@Override
public void onSuccess(Response<BatchKVResponse<Long, UpdateStatus>> response) {
results.putAll(response.getEntity().getResults());
synchronized (responses) {
responses.add(response.getEntity());
}
latch.countDown();
}
@Override
public void onError(Throwable e) {
synchronized (errors) {
errors.add(e);
}
latch.countDown();
}
};
BatchRequest<BatchKVResponse<Long, UpdateStatus>> request = requestInfo.getRequest();
RequestContext requestContext = requestInfo.getRequestContext();
REST_CLIENT.sendRequest(request, requestContext, cb);
}
latch.await();
if (!errors.isEmpty()) {
Assert.fail("Errors in scatter/gather: " + errors.toString());
}
Assert.assertEquals(results.values().size(), requestIds.length);
Set<Set<Long>> responseIdSets = new HashSet<Set<Long>>();
Set<Long> responseIds = new HashSet<Long>();
for (BatchKVResponse<Long, UpdateStatus> response : responses) {
Set<Long> theseIds = response.getResults().keySet();
Assert.assertFalse(responseIdSets.contains(theseIds)); // no duplicate requests
for (Long id : theseIds) {
Assert.assertFalse(responseIds.contains(id)); // no duplicate ids
responseIds.add(id);
}
responseIdSets.add(theseIds);
}
Assert.assertTrue(responseIds.containsAll(Arrays.asList(requestIds)));
Assert.assertEquals(responseIds.size(), requestIds.length);
}
@Test(
dataProvider =
com.linkedin.restli.internal.common.TestConstants.RESTLI_PROTOCOL_1_2_PREFIX
+ "clientDataBatchDataProvider")
public void testBatchGetKV(RestClient restClient, RestliRequestOptions requestOptions)
throws RemoteInvocationException {
List<Long> ids = Arrays.asList(1L, 2L, 3L, 4L);
Request<BatchKVResponse<Long, Greeting>> request =
new GreetingsBuilders(requestOptions).batchGet().ids(ids).buildKV();
Response<BatchKVResponse<Long, Greeting>> response =
restClient.sendRequest(request).getResponse();
BatchKVResponse<Long, Greeting> batchResponse = response.getEntity();
Assert.assertEquals(batchResponse.getResults().size(), ids.size());
}