@Override
  public BatchKVResponse<K, EntityResponse<V>> wrapResponse(DataMap dataMap, Map<String, String> headers, ProtocolVersion version)
    throws InstantiationException, IllegalAccessException, InvocationTargetException, NoSuchMethodException, IOException
  {
    if (dataMap == null)
    {
      return null;
    }

    return new BatchEntityResponse<>(dataMap, _keyType, _entityType, _keyParts, _complexKeyType, version);
  }
@Test(dataProvider = TestConstants.RESTLI_PROTOCOL_1_2_PREFIX + "batchEntityResponseDataProvider")
  public void testDecoding(List<String> keys, ProtocolVersion protocolVersion)
    throws InstantiationException, IllegalAccessException, InvocationTargetException, NoSuchMethodException, IOException
  {
    final String resultKey = keys.get(0);
    final String statusKey = keys.get(1);
    final String errorKey = keys.get(2);

    final DataMap resultData = new DataMap();
    resultData.put(resultKey, _record.data());

    final DataMap statusData = new DataMap();
    statusData.put(statusKey, _status.getCode());

    final DataMap errorData = new DataMap();
    errorData.put(errorKey, _error.data());

    final DataMap data = new DataMap();
    data.put(BatchResponse.RESULTS, resultData);
    data.put(BatchResponse.STATUSES, statusData);
    data.put(BatchResponse.ERRORS, errorData);

    final BatchEntityResponseDecoder<String, TestRecord> decoder =
        new BatchEntityResponseDecoder<>(new TypeSpec<>(TestRecord.class),
            new TypeSpec<>(String.class),
            Collections.<String, CompoundKey.TypeInfo>emptyMap(),
            null);

    final BatchKVResponse<String, EntityResponse<TestRecord>> response = decoder.wrapResponse(data, Collections.<String, String>emptyMap(), protocolVersion);
    final Map<String, EntityResponse<TestRecord>> results = response.getResults();
    final Map<String, ErrorResponse> errors = response.getErrors();

    final Collection<String> uniqueKeys = new HashSet<>(keys);
    Assert.assertEquals(results.size(), uniqueKeys.size());
    Assert.assertEquals(errors.size(), 1);

    Assert.assertEquals(results.get(resultKey).getEntity(), _record);
    Assert.assertEquals(results.get(statusKey).getStatus(), _status);
    Assert.assertEquals(results.get(errorKey).getError(), _error);
    Assert.assertEquals(errors.get(errorKey), _error);

    // Check that the response still contains the original data map
    Assert.assertEquals(response.data(), data);
  }