How to write unit testing for streaming API

Viewed 18

I have a consume function which get streaming response from a Restful API. How could I write JUnit testing on it?

I tried to create a test for the consume function but it looks like skipping the restTemplate execute call completely. Is there a way I could create a mock stream for the response so that I could process the data batches?

@Override
public void consume(
        final String url, final String jwt,
        final Consumer<StreamingDataBatch> dataBatchConsumer) {

    final StreamingQueryRequest request = StreamingParameterizedQueryRequest.builder().responseBatchSize(batchSize).build();

    restTemplate.execute(url, HttpMethod.POST,
            req -> {
                req.getHeaders().setContentType(MediaType.APPLICATION_JSON);
                req.getHeaders().set("Authorization", jwt);
                restTemplate.httpEntityCallback(request)
                        .doWithRequest(req);
            },
            resp -> {
                processStream(resp.getBody(), dataBatchConsumer);
                return null;
            });
}

private void processStream(
        final InputStream inputStream,
        final Consumer<StreamingDataBatch> dataBatchConsumer) throws IOException {

    final JsonFactory jsonFactory = objectMapper.getFactory();
    try (final JsonParser jsonParser = jsonFactory.createParser(inputStream)) {
        final JsonToken arrayToken = jsonParser.nextToken();
        if (!JsonToken.START_ARRAY.equals(arrayToken)) {
            throw new IllegalStateException("The response is not a JSON array.");
        }
        // Iterate through the objects of the array.
        while (JsonToken.START_OBJECT.equals(jsonParser.nextToken())) {
            final StreamingDataBatch batchData = jsonParser.readValueAs(
                    StreamingDataBatch.class);
            dataBatchConsumer.accept(batchData);
        }
    }
}
0 Answers
Related