org.elasticsearch.rest.RestResponseUtils Maven / Gradle / Ivy

Go to download

Show more of this group Show more artifacts with this name
Show all versions of framework Show documentation

Elasticsearch subproject :test:framework

There is a newer version: 8.16.0

/*
 * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
 * or more contributor license agreements. Licensed under the Elastic License
 * 2.0 and the Server Side Public License, v 1; you may not use this file except
 * in compliance with, at your election, the Elastic License 2.0 or the Server
 * Side Public License, v 1.
 */

package org.elasticsearch.rest;

import org.elasticsearch.common.bytes.BytesReference;
import org.elasticsearch.common.io.stream.BytesStreamOutput;
import org.elasticsearch.core.CheckedConsumer;
import org.elasticsearch.test.ESTestCase;

import java.io.IOException;
import java.io.StringWriter;
import java.io.Writer;
import java.util.Iterator;

import static org.elasticsearch.transport.BytesRefRecycler.NON_RECYCLING_INSTANCE;

public class RestResponseUtils {
    private RestResponseUtils() {}

    public static BytesReference getBodyContent(RestResponse restResponse) {
        if (restResponse.isChunked() == false) {
            return restResponse.content();
        }

        final var firstResponseBodyPart = restResponse.chunkedContent();
        assert firstResponseBodyPart.isPartComplete() == false;

        final int pageSize;
        try (var page = NON_RECYCLING_INSTANCE.obtain()) {
            pageSize = page.v().length;
        }

        try (var out = new BytesStreamOutput()) {
            while (firstResponseBodyPart.isPartComplete() == false) {
                try (var chunk = firstResponseBodyPart.encodeChunk(pageSize, NON_RECYCLING_INSTANCE)) {
                    chunk.writeTo(out);
                }
            }
            assert firstResponseBodyPart.isLastPart() : "RestResponseUtils#getBodyContent does not support continuations (yet)";

            out.flush();
            return out.bytes();
        } catch (Exception e) {
            return ESTestCase.fail(e);
        }
    }

    public static String getTextBodyContent(Iterator> iterator) {
        try (var writer = new StringWriter()) {
            while (iterator.hasNext()) {
                iterator.next().accept(writer);
            }
            writer.flush();
            return writer.toString();
        } catch (Exception e) {
            return ESTestCase.fail(e);
        }
    }
}