Search in sources :

Example 1 with VariableByteStringWriter

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

the class TestMIMEIntegrationReaderDrain method executeRequestWithDrainStrategy.

///////////////////////////////////////////////////////////////////////////////////////
private MimeMultipart executeRequestWithDrainStrategy(final int chunkSize, final List<MimeBodyPart> bodyPartList, final String drainStrategy, final String serverHeaderPrefix) throws Exception {
    MimeMultipart multiPartMimeBody = new MimeMultipart();
    //Add your body parts
    for (final MimeBodyPart bodyPart : bodyPartList) {
        multiPartMimeBody.addBodyPart(bodyPart);
    }
    final ByteArrayOutputStream byteArrayOutputStream = new ByteArrayOutputStream();
    multiPartMimeBody.writeTo(byteArrayOutputStream);
    final ByteString requestPayload = ByteString.copy(byteArrayOutputStream.toByteArray());
    final VariableByteStringWriter variableByteStringWriter = new VariableByteStringWriter(requestPayload, chunkSize);
    final EntityStream entityStream = EntityStreams.newEntityStream(variableByteStringWriter);
    final StreamRequestBuilder builder = new StreamRequestBuilder(Bootstrap.createHttpURI(PORT, SERVER_URI));
    StreamRequest request = builder.setMethod("POST").setHeader(HEADER_CONTENT_TYPE, multiPartMimeBody.getContentType()).setHeader(DRAIN_HEADER, drainStrategy).build(entityStream);
    final AtomicInteger status = new AtomicInteger(-1);
    final CountDownLatch latch = new CountDownLatch(1);
    final Map<String, String> responseHeaders = new HashMap<String, String>();
    Callback<StreamResponse> callback = expectSuccessCallback(latch, status, responseHeaders);
    _client.streamRequest(request, callback);
    latch.await(TEST_TIMEOUT, TimeUnit.MILLISECONDS);
    Assert.assertEquals(status.get(), RestStatus.OK);
    Assert.assertEquals(responseHeaders.get(DRAIN_HEADER), serverHeaderPrefix + drainStrategy);
    return multiPartMimeBody;
}
Also used : VariableByteStringWriter(com.linkedin.multipart.utils.VariableByteStringWriter) HashMap(java.util.HashMap) ByteString(com.linkedin.data.ByteString) StreamResponse(com.linkedin.r2.message.stream.StreamResponse) ByteArrayOutputStream(java.io.ByteArrayOutputStream) ByteString(com.linkedin.data.ByteString) CountDownLatch(java.util.concurrent.CountDownLatch) StreamRequestBuilder(com.linkedin.r2.message.stream.StreamRequestBuilder) StreamRequest(com.linkedin.r2.message.stream.StreamRequest) EntityStream(com.linkedin.r2.message.stream.entitystream.EntityStream) MimeMultipart(javax.mail.internet.MimeMultipart) AtomicInteger(java.util.concurrent.atomic.AtomicInteger) MimeBodyPart(javax.mail.internet.MimeBodyPart)

Example 2 with VariableByteStringWriter

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

the class TestMIMEIntegrationReader method executeRequestAndAssert.

///////////////////////////////////////////////////////////////////////////////////////
private void executeRequestAndAssert(final ByteString requestPayload, final int chunkSize, final MimeMultipart mimeMultipart) throws Exception {
    final VariableByteStringWriter variableByteStringWriter = new VariableByteStringWriter(requestPayload, chunkSize);
    final EntityStream entityStream = EntityStreams.newEntityStream(variableByteStringWriter);
    final StreamRequestBuilder builder = new StreamRequestBuilder(Bootstrap.createHttpURI(PORT, SERVER_URI));
    //We add additional parameters since MIME supports this and we want to make sure we can still extract boundary
    //properly.
    final String contentTypeHeader = mimeMultipart.getContentType() + ";somecustomparameter=somecustomvalue" + ";anothercustomparameter=anothercustomvalue";
    StreamRequest request = builder.setMethod("POST").setHeader(HEADER_CONTENT_TYPE, contentTypeHeader).build(entityStream);
    final AtomicInteger status = new AtomicInteger(-1);
    final CountDownLatch latch = new CountDownLatch(1);
    Callback<StreamResponse> callback = expectSuccessCallback(latch, status, new HashMap<String, String>());
    _client.streamRequest(request, callback);
    latch.await(TEST_TIMEOUT, TimeUnit.MILLISECONDS);
    Assert.assertEquals(status.get(), RestStatus.OK);
    List<SinglePartMIMEReaderCallbackImpl> singlePartMIMEReaderCallbacks = _mimeServerRequestHandler._testMultiPartMIMEReaderCallback._singlePartMIMEReaderCallbacks;
    Assert.assertEquals(singlePartMIMEReaderCallbacks.size(), mimeMultipart.getCount());
    for (int i = 0; i < singlePartMIMEReaderCallbacks.size(); i++) {
        //Actual
        final SinglePartMIMEReaderCallbackImpl currentCallback = singlePartMIMEReaderCallbacks.get(i);
        //Expected
        final BodyPart currentExpectedPart = mimeMultipart.getBodyPart(i);
        //Construct expected headers and verify they match
        final Map<String, String> expectedHeaders = new HashMap<String, String>();
        @SuppressWarnings("unchecked") final Enumeration<Header> allHeaders = currentExpectedPart.getAllHeaders();
        while (allHeaders.hasMoreElements()) {
            final Header header = allHeaders.nextElement();
            expectedHeaders.put(header.getName(), header.getValue());
        }
        Assert.assertEquals(currentCallback._headers, expectedHeaders);
        //Verify the body matches
        if (currentExpectedPart.getContent() instanceof byte[]) {
            Assert.assertEquals(currentCallback._finishedData.copyBytes(), currentExpectedPart.getContent());
        } else {
            //Default is String
            Assert.assertEquals(new String(currentCallback._finishedData.copyBytes()), currentExpectedPart.getContent());
        }
    }
}
Also used : VariableByteStringWriter(com.linkedin.multipart.utils.VariableByteStringWriter) MimeBodyPart(javax.mail.internet.MimeBodyPart) BodyPart(javax.mail.BodyPart) HashMap(java.util.HashMap) StreamResponse(com.linkedin.r2.message.stream.StreamResponse) ByteString(com.linkedin.data.ByteString) CountDownLatch(java.util.concurrent.CountDownLatch) StreamRequestBuilder(com.linkedin.r2.message.stream.StreamRequestBuilder) StreamRequest(com.linkedin.r2.message.stream.StreamRequest) EntityStream(com.linkedin.r2.message.stream.entitystream.EntityStream) Header(javax.mail.Header) AtomicInteger(java.util.concurrent.atomic.AtomicInteger)

Aggregations

ByteString (com.linkedin.data.ByteString)2 VariableByteStringWriter (com.linkedin.multipart.utils.VariableByteStringWriter)2 StreamRequest (com.linkedin.r2.message.stream.StreamRequest)2 StreamRequestBuilder (com.linkedin.r2.message.stream.StreamRequestBuilder)2 StreamResponse (com.linkedin.r2.message.stream.StreamResponse)2 EntityStream (com.linkedin.r2.message.stream.entitystream.EntityStream)2 HashMap (java.util.HashMap)2 CountDownLatch (java.util.concurrent.CountDownLatch)2 AtomicInteger (java.util.concurrent.atomic.AtomicInteger)2 MimeBodyPart (javax.mail.internet.MimeBodyPart)2 ByteArrayOutputStream (java.io.ByteArrayOutputStream)1 BodyPart (javax.mail.BodyPart)1 Header (javax.mail.Header)1 MimeMultipart (javax.mail.internet.MimeMultipart)1