Beispiel #1
0
  @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());
  }