Examples with MultiPartMIMEFullReaderCallback com.linkedin.multipart.utils.MIMETestUtils.MultiPartMIMEFullReaderCallback used on opensource projects

Example 1 with MultiPartMIMEFullReaderCallback

use of com.linkedin.multipart.utils.MIMETestUtils.MultiPartMIMEFullReaderCallback 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 2 with MultiPartMIMEFullReaderCallback

use of com.linkedin.multipart.utils.MIMETestUtils.MultiPartMIMEFullReaderCallback in project rest.li by linkedin.

the class TestRestLiServer method testRequestAttachmentsAndResponseAttachments.

@Test
public void testRequestAttachmentsAndResponseAttachments() throws Exception {
    // This test verifies the server's ability to accept request attachments and send back response attachments. This is the
    // main test to verify the wire protocol for streaming. We send a payload that contains the rest.li payload and some attachments
    // and we send a response back with a rest.li payload and some attachments.
    // Define the server side resource attachments to be sent back.
    final RestLiResponseAttachments.Builder responseAttachmentsBuilder = new RestLiResponseAttachments.Builder();
    responseAttachmentsBuilder.appendSingleAttachment(new RestLiTestAttachmentDataSource("1", ByteString.copyString("one", Charset.defaultCharset())));
    Capture<ResourceContext> resourceContextCapture = EasyMock.newCapture();
    final AsyncStatusCollectionResource statusResource = getMockResource(AsyncStatusCollectionResource.class, capture(resourceContextCapture));
    statusResource.streamingAction(EasyMock.<String>anyObject(), EasyMock.<RestLiAttachmentReader>anyObject(), EasyMock.<Callback<Long>>anyObject());
    EasyMock.expectLastCall().andAnswer(new IAnswer<Object>() {

        @Override
        public Object answer() throws Throwable {
            // Verify there are still attachments to be read.
            final RestLiAttachmentReader attachmentReader = (RestLiAttachmentReader) EasyMock.getCurrentArguments()[1];
            Assert.assertFalse(attachmentReader.haveAllAttachmentsFinished());
            // Verify the action param.
            Assert.assertEquals((String) EasyMock.getCurrentArguments()[0], "someMetadata");
            // Set the response attachments
            resourceContextCapture.getValue().setResponseAttachments(responseAttachmentsBuilder.build());
            // Now respond back to the request.
            @SuppressWarnings("unchecked") Callback<Long> callback = (Callback<Long>) EasyMock.getCurrentArguments()[2];
            callback.onSuccess(1234l);
            return null;
        }
    });
    replay(statusResource);
    // Now we create a multipart/related payload.
    final String payload = "{\"metadata\": \"someMetadata\"}";
    final ByteStringWriter byteStringWriter = new ByteStringWriter(ByteString.copyString(payload, Charset.defaultCharset()));
    final MultiPartMIMEWriter.Builder builder = new MultiPartMIMEWriter.Builder();
    AttachmentUtils.appendSingleAttachmentToBuilder(builder, new RestLiTestAttachmentDataSource("2", ByteString.copyString("two", Charset.defaultCharset())));
    final MultiPartMIMEWriter writer = AttachmentUtils.createMultiPartMIMEWriter(byteStringWriter, "application/json", builder);
    final StreamRequest streamRequest = MultiPartMIMEStreamRequestFactory.generateMultiPartMIMEStreamRequest(new URI("/asyncstatuses/?action=streamingAction"), "related", writer, Collections.<String, String>emptyMap(), "POST", ImmutableMap.of(RestConstants.HEADER_ACCEPT, RestConstants.HEADER_VALUE_MULTIPART_RELATED), Collections.emptyList());
    final Callback<StreamResponse> callback = new Callback<StreamResponse>() {

        @Override
        public void onSuccess(StreamResponse streamResponse) {
            // Before reading the data make sure top level type is multipart/related
            Assert.assertEquals(streamResponse.getStatus(), 200);
            try {
                ContentType contentType = new ContentType(streamResponse.getHeader(RestConstants.HEADER_CONTENT_TYPE));
                Assert.assertEquals(contentType.getBaseType(), RestConstants.HEADER_VALUE_MULTIPART_RELATED);
            } catch (ParseException parseException) {
                Assert.fail();
            }
            final CountDownLatch countDownLatch = new CountDownLatch(1);
            MultiPartMIMEFullReaderCallback fullReaderCallback = new MultiPartMIMEFullReaderCallback(countDownLatch);
            final MultiPartMIMEReader reader = MultiPartMIMEReader.createAndAcquireStream(streamResponse);
            reader.registerReaderCallback(fullReaderCallback);
            try {
                countDownLatch.await(3000, TimeUnit.MILLISECONDS);
            } catch (InterruptedException interruptedException) {
                Assert.fail();
            }
            final List<SinglePartMIMEFullReaderCallback> singlePartMIMEReaderCallbacks = fullReaderCallback.getSinglePartMIMEReaderCallbacks();
            Assert.assertEquals(singlePartMIMEReaderCallbacks.size(), 2);
            // Verify first part is Action response.
            Assert.assertEquals(singlePartMIMEReaderCallbacks.get(0).getHeaders().get(RestConstants.HEADER_CONTENT_TYPE), RestConstants.HEADER_VALUE_APPLICATION_JSON);
            Assert.assertEquals(singlePartMIMEReaderCallbacks.get(0).getFinishedData().asAvroString(), "{\"value\":1234}");
            // Verify the second part matches what the server should have sent back
            Assert.assertEquals(singlePartMIMEReaderCallbacks.get(1).getHeaders().get(RestConstants.HEADER_CONTENT_ID), "1");
            Assert.assertEquals(singlePartMIMEReaderCallbacks.get(1).getFinishedData().asString(Charset.defaultCharset()), "one");
            EasyMock.verify(statusResource);
            EasyMock.reset(statusResource);
        }

        @Override
        public void onError(Throwable e) {
            fail();
        }
    };
    _server.handleRequest(streamRequest, new RequestContext(), callback);
}