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

Example 1 with MultiPartMIMEReaderCallback

use of com.linkedin.multipart.MultiPartMIMEReaderCallback 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());
}