@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 + "protocol")
  public void testGatherBatchResponse(ProtocolVersion version)
  {
    Map<RequestInfo, Response<BatchResponse<TestRecord>>> successResponses = new HashMap<>();
    successResponses.put(
            new RequestInfo(createBatchGetRequest(1L, 2L), getTargetHostRequestContext(_host1URI)),
            createBatchResponse(Collections.singleton(1L), Collections.singleton(2L)));
    Map<RequestInfo, Throwable> failResponses = new HashMap<>();
    failResponses.put(
            new RequestInfo(createBatchGetRequest(3L), getTargetHostRequestContext(_host2URI)),
            new RestLiScatterGatherException("Partition host is unavailable!"));
    Callback<Response<BatchResponse<TestRecord>>> testCallback = new Callback<Response<BatchResponse<TestRecord>>>()
    {
      @Override
      public void onError(Throwable e)
      {

      }

      @Override
      public void onSuccess(Response<BatchResponse<TestRecord>> result)
      {
        Assert.assertNotNull(result.getEntity());
        Assert.assertEquals(result.getStatus(), HttpStatus.S_200_OK.getCode());
        Assert.assertTrue(result.getEntity().getResults().size() == 1);
        Assert.assertTrue(result.getEntity().getResults().containsKey("1"));
        Assert.assertTrue(result.getEntity().getErrors().size() == 3);
        ErrorResponse keyError = result.getEntity().getErrors().get("2");
        Assert.assertEquals(keyError.getStatus().intValue(), HttpStatus.S_404_NOT_FOUND.getCode());
        ErrorResponse failError = result.getEntity().getErrors().get("3");
        Assert.assertEquals(failError.getExceptionClass(), RestLiScatterGatherException.class.getName());
        Assert.assertEquals(failError.getMessage(), "Partition host is unavailable!");
        ErrorResponse unmappedError = result.getEntity().getErrors().get("4");
        Assert.assertEquals(unmappedError.getExceptionClass(), RestLiScatterGatherException.class.getName());
        Assert.assertEquals(unmappedError.getMessage(), "Unable to find a host for keys :[4]");
      }
    };
    _sgStrategy.onAllResponsesReceived(_batchGetRequest, version, successResponses, failResponses, _unmappedKeys, testCallback);
  }