@SuppressWarnings({"unchecked", "rawtypes"})
  public void testTooLargeResponse() throws Exception {
    ContentTooLongException tooLong = new ContentTooLongException("too long!");
    CloseableHttpAsyncClient httpClient = mock(CloseableHttpAsyncClient.class);
    when(httpClient.<HttpResponse>execute(
            any(HttpAsyncRequestProducer.class),
            any(HttpAsyncResponseConsumer.class),
            any(FutureCallback.class)))
        .then(
            new Answer<Future<HttpResponse>>() {
              @Override
              public Future<HttpResponse> answer(InvocationOnMock invocationOnMock)
                  throws Throwable {
                HeapBufferedAsyncResponseConsumer consumer =
                    (HeapBufferedAsyncResponseConsumer) invocationOnMock.getArguments()[1];
                FutureCallback callback = (FutureCallback) invocationOnMock.getArguments()[2];

                assertEquals(
                    new ByteSizeValue(200, ByteSizeUnit.MB).bytesAsInt(),
                    consumer.getBufferLimit());
                callback.failed(tooLong);
                return null;
              }
            });
    RemoteScrollableHitSource source = sourceWithMockedClient(true, httpClient);

    AtomicBoolean called = new AtomicBoolean();
    Consumer<Response> checkResponse = r -> called.set(true);
    Throwable e =
        expectThrows(
            RuntimeException.class,
            () -> source.doStartNextScroll(FAKE_SCROLL_ID, timeValueMillis(0), checkResponse));
    // Unwrap the some artifacts from the test
    while (e.getMessage().equals("failed")) {
      e = e.getCause();
    }
    // This next exception is what the user sees
    assertEquals(
        "Remote responded with a chunk that was too large. Use a smaller batch size.",
        e.getMessage());
    // And that exception is reported as being caused by the underlying exception returned by the
    // client
    assertSame(tooLong, e.getCause());
    assertFalse(called.get());
  }
  public void testWrapExceptionToPreserveStatus() throws IOException {
    Exception cause = new Exception();

    // Successfully get the status without a body
    RestStatus status = randomFrom(RestStatus.values());
    ElasticsearchStatusException wrapped =
        RemoteScrollableHitSource.wrapExceptionToPreserveStatus(status.getStatus(), null, cause);
    assertEquals(status, wrapped.status());
    assertEquals(cause, wrapped.getCause());
    assertEquals("No error body.", wrapped.getMessage());

    // Successfully get the status without a body
    HttpEntity okEntity = new StringEntity("test body", StandardCharsets.UTF_8);
    wrapped =
        RemoteScrollableHitSource.wrapExceptionToPreserveStatus(
            status.getStatus(), okEntity, cause);
    assertEquals(status, wrapped.status());
    assertEquals(cause, wrapped.getCause());
    assertEquals("body=test body", wrapped.getMessage());

    // Successfully get the status with a broken body
    IOException badEntityException = new IOException();
    HttpEntity badEntity = mock(HttpEntity.class);
    when(badEntity.getContent()).thenThrow(badEntityException);
    wrapped =
        RemoteScrollableHitSource.wrapExceptionToPreserveStatus(
            status.getStatus(), badEntity, cause);
    assertEquals(status, wrapped.status());
    assertEquals(cause, wrapped.getCause());
    assertEquals("Failed to extract body.", wrapped.getMessage());
    assertEquals(badEntityException, wrapped.getSuppressed()[0]);

    // Fail to get the status without a body
    int notAnHttpStatus = -1;
    assertNull(RestStatus.fromCode(notAnHttpStatus));
    wrapped = RemoteScrollableHitSource.wrapExceptionToPreserveStatus(notAnHttpStatus, null, cause);
    assertEquals(RestStatus.INTERNAL_SERVER_ERROR, wrapped.status());
    assertEquals(cause, wrapped.getCause());
    assertEquals(
        "Couldn't extract status [" + notAnHttpStatus + "]. No error body.", wrapped.getMessage());

    // Fail to get the status without a body
    wrapped =
        RemoteScrollableHitSource.wrapExceptionToPreserveStatus(notAnHttpStatus, okEntity, cause);
    assertEquals(RestStatus.INTERNAL_SERVER_ERROR, wrapped.status());
    assertEquals(cause, wrapped.getCause());
    assertEquals(
        "Couldn't extract status [" + notAnHttpStatus + "]. body=test body", wrapped.getMessage());

    // Fail to get the status with a broken body
    wrapped =
        RemoteScrollableHitSource.wrapExceptionToPreserveStatus(notAnHttpStatus, badEntity, cause);
    assertEquals(RestStatus.INTERNAL_SERVER_ERROR, wrapped.status());
    assertEquals(cause, wrapped.getCause());
    assertEquals(
        "Couldn't extract status [" + notAnHttpStatus + "]. Failed to extract body.",
        wrapped.getMessage());
    assertEquals(badEntityException, wrapped.getSuppressed()[0]);
  }