Examples with MultiPartMIMEReader com.linkedin.multipart.MultiPartMIMEReader used on opensource projects

Example 6 with MultiPartMIMEReader

use of com.linkedin.multipart.MultiPartMIMEReader in project rest.li by linkedin.

the class TestParseqTraceDebugRequestHandler method testErrorStreamingAbsorbRequestAbortResponse.

@Test
public void testErrorStreamingAbsorbRequestAbortResponse() throws Exception {
    //This test verifies that in the face of an error, the ParseqTraceDebugRequestHandler aborts any potential outgoing
    //response attachments and absorbs and drops on the ground any incoming request attachments.
    final URI uri = URI.create("http://host/abc/12/__debug/parseqtrace/raw");
    ParseqTraceDebugRequestHandler requestHandler = new ParseqTraceDebugRequestHandler();
    RestRequestBuilder requestBuilder = new RestRequestBuilder(uri);
    RestRequest request = requestBuilder.build();
    RequestContext requestContext = new RequestContext();
    final RequestExecutionCallback<RestResponse> callback = new RequestExecutionCallback<RestResponse>() {

        @Override
        public void onError(Throwable e, RequestExecutionReport executionReport, RestLiAttachmentReader requestAttachmentReader, RestLiResponseAttachments responseAttachments) {
            //Even though we call callback.onError() below to simulate a failed rest.li request execution, the
            //ParseqTraceDebugRequestHandler eventually treats this as a success.
            Assert.fail("Request execution passed unexpectedly.");
        }

        @Override
        public void onSuccess(RestResponse result, RequestExecutionReport executionReport, RestLiResponseAttachments responseAttachments) {
        }
    };
    //Create request and response attachments. This is not the wire protocol for rest.li streaming, but
    //makes this test easier to write and understand.
    //Response attachment.
    final RestLiTestAttachmentDataSource responseDataSource = RestLiTestAttachmentDataSource.createWithRandomPayload("1");
    final RestLiResponseAttachments responseAttachments = new RestLiResponseAttachments.Builder().appendSingleAttachment(responseDataSource).build();
    //Request attachment. We need to create a reader here.
    final RestLiTestAttachmentDataSource requestAttachment = RestLiTestAttachmentDataSource.createWithRandomPayload("2");
    final MultiPartMIMEWriter.Builder builder = new MultiPartMIMEWriter.Builder();
    AttachmentUtils.appendSingleAttachmentToBuilder(builder, requestAttachment);
    final ByteStringWriter byteStringWriter = new ByteStringWriter(ByteString.copyString("Normal rest.li request payload", Charset.defaultCharset()));
    final MultiPartMIMEWriter writer = AttachmentUtils.createMultiPartMIMEWriter(byteStringWriter, "application/json", builder);
    final StreamRequest streamRequest = MultiPartMIMEStreamRequestFactory.generateMultiPartMIMEStreamRequest(new URI(""), "related", writer);
    final MultiPartMIMEReader reader = MultiPartMIMEReader.createAndAcquireStream(streamRequest);
    final RestLiAttachmentReader attachmentReader = new RestLiAttachmentReader(reader);
    //Absorb the first part as rest.li server does and then give the beginning of the next part to the
    //ParseqTraceDebugRequestHandler.
    final MultiPartMIMEReaderCallback multiPartMIMEReaderCallback = new MultiPartMIMEReaderCallback() {

        int partCounter = 0;

        @Override
        public void onNewPart(MultiPartMIMEReader.SinglePartMIMEReader singlePartMIMEReader) {
            if (partCounter == 0) {
                singlePartMIMEReader.drainPart();
                partCounter++;
            } else {
                //When the first part is read in, at the beginning of the 2nd part, we move to the debug request handler.
                requestHandler.handleRequest(request, requestContext, new RestLiDebugRequestHandler.ResourceDebugRequestHandler() {

                    @Override
                    public void handleRequest(RestRequest request, RequestContext requestContext, RequestExecutionCallback<RestResponse> callback) {
                        callback.onError(RestException.forError(500, "An error has occurred"), new RequestExecutionReportBuilder().build(), attachmentReader, responseAttachments);
                    }
                }, attachmentReader, callback);
            }
        }

        @Override
        public void onFinished() {
            Assert.fail();
        }

        @Override
        public void onDrainComplete() {
        //Eventually this should occur.
        }

        @Override
        public void onStreamError(Throwable throwable) {
            Assert.fail();
        }
    };
    reader.registerReaderCallback(multiPartMIMEReaderCallback);
    //The response data source should have been aborted.
    Assert.assertTrue(responseDataSource.dataSourceAborted());
    //The request attachment should have been absorbed and finished.
    Assert.assertTrue(requestAttachment.finished());
}

Example 7 with MultiPartMIMEReader

use of com.linkedin.multipart.MultiPartMIMEReader in project rest.li by linkedin.

the class TestRestClientRequestBuilder method verifyStreamRequest.

private void verifyStreamRequest(final StreamRequest streamRequest, final boolean acceptResponseAttachments, final String expectedAcceptHeader, final boolean streamAttachments, final String expectedContentTypeHeader, final String expectedRequestBody) {
    // Otherwise it went down the RestRequest code path incorrectly.
    Assert.assertNotNull(streamRequest);
    // The accept type header will look different based on whether or not attachments were expected.
    if (acceptResponseAttachments) {
        if (expectedAcceptHeader != null) {
            Assert.assertTrue(streamRequest.getHeader(ACCEPT_TYPE_HEADER).startsWith(expectedAcceptHeader));
            Assert.assertTrue(streamRequest.getHeader(ACCEPT_TYPE_HEADER).contains(RestConstants.HEADER_VALUE_MULTIPART_RELATED));
        } else {
            Assert.assertEquals(streamRequest.getHeader(ACCEPT_TYPE_HEADER), RestConstants.HEADER_VALUE_MULTIPART_RELATED + ";q=1.0");
        }
    } else {
        Assert.assertEquals(streamRequest.getHeader(ACCEPT_TYPE_HEADER), expectedAcceptHeader);
    }
    if (!streamAttachments) {
        // Verify content type header.
        Assert.assertEquals(streamRequest.getHeader(CONTENT_TYPE_HEADER), expectedContentTypeHeader);
        // If there are no attachments, then we can just read everything in
        Messages.toRestRequest(streamRequest, new Callback<RestRequest>() {

            @Override
            public void onError(Throwable e) {
                Assert.fail();
            }

            @Override
            public void onSuccess(RestRequest result) {
                // Verify entity after the conversion is complete.
                Assert.assertEquals(result.getEntity().asAvroString(), expectedRequestBody);
            }
        });
    } else {
        // There were attachments so let's read using MultiPartMIMEReader to verify the wire format designed by RestClient
        // is indeed correct
        final MultiPartMIMEReader streamRequestReader = MultiPartMIMEReader.createAndAcquireStream(streamRequest);
        final CountDownLatch streamRequestReaderLatch = new CountDownLatch(1);
        final MultiPartMIMEFullReaderCallback streamRequestReaderCallback = new MultiPartMIMEFullReaderCallback(streamRequestReaderLatch);
        streamRequestReader.registerReaderCallback(streamRequestReaderCallback);
        try {
            streamRequestReaderLatch.await(3000, TimeUnit.MILLISECONDS);
        } catch (InterruptedException interruptedException) {
            Assert.fail();
        }
        final List<SinglePartMIMEFullReaderCallback> singlePartMIMEReaderCallbacks = streamRequestReaderCallback.getSinglePartMIMEReaderCallbacks();
        // We have should have three parts. One for the rest.li payload and two for the attachments.
        Assert.assertEquals(singlePartMIMEReaderCallbacks.size(), 3);
        // Verify the first part by looking at its content type and payload.
        Assert.assertEquals(singlePartMIMEReaderCallbacks.get(0).getHeaders().get(CONTENT_TYPE_HEADER), expectedContentTypeHeader);
        Assert.assertEquals(singlePartMIMEReaderCallbacks.get(0).getFinishedData().asAvroString(), expectedRequestBody);
        // Verify the top level content type is multipart mime related. Use startsWith() since the boundary is random.
        Assert.assertTrue(streamRequest.getHeader(CONTENT_TYPE_HEADER).startsWith(RestConstants.HEADER_VALUE_MULTIPART_RELATED));
        // Now verify the attachments. We have to remove the first part since we already read it.
        singlePartMIMEReaderCallbacks.remove(0);
        verifyAttachments(singlePartMIMEReaderCallbacks);
    }
}

Example 8 with MultiPartMIMEReader

use of com.linkedin.multipart.MultiPartMIMEReader in project rest.li by linkedin.

the class TestAttachmentUtils method testMixtureOfAttachments.

@Test
public void testMixtureOfAttachments() {
    final MultiPartMIMEWriter.Builder builder = new MultiPartMIMEWriter.Builder();
    final List<RestLiTestAttachmentDataSource> iteratorDataSources = generateTestDataSources();
    final RestLiTestAttachmentDataSourceIterator dataSourceIterator = new RestLiTestAttachmentDataSourceIterator(iteratorDataSources, new IllegalArgumentException());
    // Let each data source know its parent, so that when the data source is done, it can notify it's parent to call onNewDataSourceWriter()
    for (final RestLiTestAttachmentDataSource dataSource : iteratorDataSources) {
        dataSource.setParentDataSourceIterator(dataSourceIterator);
    }
    // Now one at the beginning
    final RestLiTestAttachmentDataSource dataSourceBeginning = RestLiTestAttachmentDataSource.createWithRandomPayload("BEG");
    // And one at the end
    final RestLiTestAttachmentDataSource dataSourceEnd = RestLiTestAttachmentDataSource.createWithRandomPayload("END");
    AttachmentUtils.appendSingleAttachmentToBuilder(builder, dataSourceBeginning);
    AttachmentUtils.appendMultipleAttachmentsToBuilder(builder, dataSourceIterator);
    AttachmentUtils.appendSingleAttachmentToBuilder(builder, dataSourceEnd);
    final StreamRequest streamRequest = MultiPartMIMEStreamRequestFactory.generateMultiPartMIMEStreamRequest(URI.create("foo"), "related", builder.build());
    final MultiPartMIMEReader streamRequestReader = MultiPartMIMEReader.createAndAcquireStream(streamRequest);
    final CountDownLatch streamRequestReaderLatch = new CountDownLatch(1);
    final MIMETestUtils.MultiPartMIMEFullReaderCallback streamRequestReaderCallback = new MIMETestUtils.MultiPartMIMEFullReaderCallback(streamRequestReaderLatch);
    streamRequestReader.registerReaderCallback(streamRequestReaderCallback);
    try {
        streamRequestReaderLatch.await(3000, TimeUnit.MILLISECONDS);
    } catch (InterruptedException interruptedException) {
        Assert.fail();
    }
    final List<RestLiTestAttachmentDataSource> allAttachmentDataSources = new ArrayList<>();
    allAttachmentDataSources.add(dataSourceBeginning);
    allAttachmentDataSources.addAll(iteratorDataSources);
    allAttachmentDataSources.add(dataSourceEnd);
    verifyAttachments(streamRequestReaderCallback.getSinglePartMIMEReaderCallbacks(), allAttachmentDataSources);
}

Example 9 with MultiPartMIMEReader

use of com.linkedin.multipart.MultiPartMIMEReader in project rest.li by linkedin.

the class TestAttachmentUtils method testSingleAttachment.

@Test
public void testSingleAttachment() {
    final MultiPartMIMEWriter.Builder builder = new MultiPartMIMEWriter.Builder();
    final List<RestLiTestAttachmentDataSource> testAttachmentDataSources = generateTestDataSources();
    for (RestLiTestAttachmentDataSource dataSource : testAttachmentDataSources) {
        AttachmentUtils.appendSingleAttachmentToBuilder(builder, dataSource);
    }
    final StreamRequest streamRequest = MultiPartMIMEStreamRequestFactory.generateMultiPartMIMEStreamRequest(URI.create("foo"), "related", builder.build());
    final MultiPartMIMEReader streamRequestReader = MultiPartMIMEReader.createAndAcquireStream(streamRequest);
    final CountDownLatch streamRequestReaderLatch = new CountDownLatch(1);
    final MIMETestUtils.MultiPartMIMEFullReaderCallback streamRequestReaderCallback = new MIMETestUtils.MultiPartMIMEFullReaderCallback(streamRequestReaderLatch);
    streamRequestReader.registerReaderCallback(streamRequestReaderCallback);
    try {
        streamRequestReaderLatch.await(3000, TimeUnit.MILLISECONDS);
    } catch (InterruptedException interruptedException) {
        Assert.fail();
    }
    verifyAttachments(streamRequestReaderCallback.getSinglePartMIMEReaderCallbacks(), testAttachmentDataSources);
}

Example 10 with MultiPartMIMEReader

use of com.linkedin.multipart.MultiPartMIMEReader in project rest.li by linkedin.

the class RestResponseDecoder method decodeResponse.

public void decodeResponse(final StreamResponse streamResponse, final Callback<Response<T>> responseCallback) throws RestLiDecodingException {
    // Determine content type and take appropriate action.
    // If 'multipart/related', then use MultiPartMIMEReader to read first part (which can be json or pson).
    final String contentTypeString = streamResponse.getHeader(RestConstants.HEADER_CONTENT_TYPE);
    if (contentTypeString != null) {
        ContentType contentType = null;
        try {
            contentType = new ContentType(contentTypeString);
        } catch (ParseException parseException) {
            responseCallback.onError(new RestLiDecodingException("Could not decode Content-Type header in response", parseException));
            return;
        }
        if (contentType.getBaseType().equalsIgnoreCase(RestConstants.HEADER_VALUE_MULTIPART_RELATED)) {
            final MultiPartMIMEReader multiPartMIMEReader = MultiPartMIMEReader.createAndAcquireStream(streamResponse);
            final TopLevelReaderCallback topLevelReaderCallback = new TopLevelReaderCallback(responseCallback, streamResponse, multiPartMIMEReader);
            multiPartMIMEReader.registerReaderCallback(topLevelReaderCallback);
            return;
        }
    }
    // Otherwise if the whole body is json/pson then read everything in.
    StreamDataCodec streamDataCodec = null;
    try {
        streamDataCodec = getContentType(streamResponse.getHeaders().get(RestConstants.HEADER_CONTENT_TYPE)).orElse(JSON).getStreamCodec();
    } catch (MimeTypeParseException e) {
        responseCallback.onError(e);
        return;
    }
    if (streamDataCodec != null) {
        CompletionStage<DataMap> dataMapCompletionStage = streamDataCodec.decodeMap(EntityStreamAdapters.toGenericEntityStream(streamResponse.getEntityStream()));
        dataMapCompletionStage.handle((dataMap, e) -> {
            if (e != null) {
                responseCallback.onError(new RestLiDecodingException("Could not decode REST response", e));
                return null;
            }
            try {
                responseCallback.onSuccess(createResponse(streamResponse.getHeaders(), streamResponse.getStatus(), dataMap, streamResponse.getCookies()));
            } catch (Throwable throwable) {
                responseCallback.onError(throwable);
            }
            // handle function requires a return statement although there is no more completion stage.
            return null;
        });
    } else {
        final FullEntityReader fullEntityReader = new FullEntityReader(new Callback<ByteString>() {

            @Override
            public void onError(Throwable e) {
                responseCallback.onError(e);
            }

            @Override
            public void onSuccess(ByteString result) {
                try {
                    responseCallback.onSuccess(createResponse(streamResponse.getHeaders(), streamResponse.getStatus(), result, streamResponse.getCookies()));
                } catch (Exception exception) {
                    onError(exception);
                }
            }
        });
        streamResponse.getEntityStream().setReader(fullEntityReader);
    }
}