@Override
  @SuppressWarnings("unchecked")
  public <K, T> void onAllResponsesReceived(Request<T> request, ProtocolVersion protocolVersion,
                                            Map<RequestInfo, Response<T>> successResponses,
                                            Map<RequestInfo, Throwable> failureResponses,
                                            Map<Integer, Set<K>> unmappedKeys,
                                            Callback<Response<T>> callback)
  {
    BatchRequest<T> batchRequest = safeCastRequest(request);
    // initialize an empty dataMap for final response entity
    DataMap gatheredResponseDataMap = initializeResponseContainer();
    // gather success response
    successResponses.forEach((req, response) -> gatherResponse(gatheredResponseDataMap, req, response.getEntity()));
    // gather failure response
    failureResponses.forEach((req, e) ->
            {
              Set<K> failedKeys = (Set<K>)((BatchRequest<T>)req.getRequest()).getObjectIds();
              gatherException(gatheredResponseDataMap, failedKeys, e, protocolVersion);
            });
    // gather unmapped keys
    if (unmappedKeys != null && !unmappedKeys.isEmpty())
    {
      Set<K> unmapped = unmappedKeys.values().stream().flatMap(Collection::stream).collect(Collectors.toSet());
      gatherException(gatheredResponseDataMap, unmapped,
              new RestLiScatterGatherException("Unable to find a host for keys :" + unmapped),
              protocolVersion);
    }
    T gatheredResponse = constructResponseFromDataMap(batchRequest, protocolVersion, gatheredResponseDataMap);
    if (!successResponses.isEmpty())
    {
      Response<T> firstResponse = successResponses.values().iterator().next();
      callback.onSuccess(new ResponseImpl<>(firstResponse, gatheredResponse));
    }
    else
    {
      // all scattered requests are failing, we still return 200 for original request, but body will contain
      // failed response for each key.
      callback.onSuccess(new ResponseImpl<>(HttpStatus.S_200_OK.getCode(),
              Collections.emptyMap(), Collections.emptyList(), gatheredResponse, null));
    }
  }
@Test(dataProvider = TestConstants.RESTLI_PROTOCOL_1_2_PREFIX + "gatherBatchResponseProvider")
  @SuppressWarnings({"rawtypes", "unchecked"})
  public void testGatherBatchKVResponse(Request request,
                                        ProtocolVersion version,
                                        Request successRequest,
                                        URI successHost,
                                        Response successResponse,
                                        Request<?> failRequest,
                                        URI failHost,
                                        int resultDataMapSize,
                                        int resultSize)
  {
    Map<RequestInfo, Response<BatchKVResponse<Long, ?>>> successResponses = new HashMap<>();
    successResponses.put(
            new RequestInfo(successRequest, getTargetHostRequestContext(successHost)),
            successResponse);
    Map<RequestInfo, Throwable> failResponses = new HashMap<>();
    failResponses.put(
            new RequestInfo(failRequest, getTargetHostRequestContext(failHost)),
            new RestLiScatterGatherException("Partition host is unavailable!"));
    Callback<Response<BatchKVResponse<Long, ?>>> testCallback =
            new Callback<Response<BatchKVResponse<Long, ?>>> ()
            {
              @Override
              public void onError(Throwable e)
              {

              }

              @Override
              public void onSuccess(Response<BatchKVResponse<Long, ?>> result)
              {
                Assert.assertNotNull(result.getEntity());
                Assert.assertEquals(result.getStatus(), HttpStatus.S_200_OK.getCode());
                Assert.assertTrue(result.getEntity().data().getDataMap(BatchResponse.RESULTS).size() == resultDataMapSize);
                // BatchKVResponse.getResults() contains all entries including both successful and failed ones.
                Assert.assertTrue(result.getEntity().getResults().size() == resultSize);
                Assert.assertTrue(result.getEntity().getResults().containsKey(1L));
                // merged error can come from 3 cases:
                // - errored keys in successfully returned scattered batch response
                // - RemoteInvocationException from a scattered request
                // - unmapped keys
                Assert.assertTrue(result.getEntity().data().getDataMap(BatchResponse.ERRORS).size() == 3);
                Assert.assertTrue(result.getEntity().getErrors().size() == 3);
                ErrorResponse keyError = result.getEntity().getErrors().get(2L);
                Assert.assertEquals(keyError.getStatus().intValue(), HttpStatus.S_404_NOT_FOUND.getCode());
                ErrorResponse failError = result.getEntity().getErrors().get(3L);
                Assert.assertEquals(failError.getExceptionClass(), RestLiScatterGatherException.class.getName());
                Assert.assertEquals(failError.getMessage(), "Partition host is unavailable!");
                ErrorResponse unmappedError = result.getEntity().getErrors().get(4L);
                Assert.assertEquals(unmappedError.getExceptionClass(), RestLiScatterGatherException.class.getName());
                Assert.assertEquals(unmappedError.getMessage(), "Unable to find a host for keys :[4]");
              }
            };
    _sgStrategy.onAllResponsesReceived(request, version, successResponses, failResponses, _unmappedKeys, testCallback);
  }