Search in sources :

Example 56 with EntityStream

use of com.linkedin.r2.message.stream.entitystream.EntityStream in project rest.li by linkedin.

the class TestMIMEReaderStateTransitions method testRegisterCallbackMultiPartMIMEReader.

// MultiPartMIMEReader exceptions:
@Test
public void testRegisterCallbackMultiPartMIMEReader() {
    final EntityStream entityStream = mock(EntityStream.class);
    final StreamRequest streamRequest = mock(StreamRequest.class);
    when(streamRequest.getEntityStream()).thenReturn(entityStream);
    when(streamRequest.getHeader(MultiPartMIMEUtils.CONTENT_TYPE_HEADER)).thenReturn("multipart/mixed; boundary=\"--123\"");
    MultiPartMIMEReader reader = MultiPartMIMEReader.createAndAcquireStream(streamRequest);
    // Test each possible exception:
    reader.setState(MultiPartMIMEReader.MultiPartReaderState.FINISHED);
    try {
        reader.registerReaderCallback(EMPTY_MULTI_PART_MIME_READER_CALLBACK);
        Assert.fail();
    } catch (MultiPartReaderFinishedException multiPartReaderFinishedException) {
    }
    reader.setState(MultiPartMIMEReader.MultiPartReaderState.READING_EPILOGUE);
    try {
        reader.registerReaderCallback(EMPTY_MULTI_PART_MIME_READER_CALLBACK);
        Assert.fail();
    } catch (MultiPartReaderFinishedException multiPartReaderFinishedException) {
    }
    reader.setState(MultiPartMIMEReader.MultiPartReaderState.CALLBACK_BOUND_AND_READING_PREAMBLE);
    try {
        reader.registerReaderCallback(EMPTY_MULTI_PART_MIME_READER_CALLBACK);
        Assert.fail();
    } catch (StreamBusyException streamBusyException) {
    }
    reader.setState(MultiPartMIMEReader.MultiPartReaderState.DRAINING);
    try {
        reader.registerReaderCallback(EMPTY_MULTI_PART_MIME_READER_CALLBACK);
        Assert.fail();
    } catch (StreamBusyException streamBusyException) {
    }
    // This is a desired top level reader state
    reader.setState(MultiPartMIMEReader.MultiPartReaderState.READING_PARTS);
    final MultiPartMIMEReader.SinglePartMIMEReader singlePartMIMEReader = reader.new SinglePartMIMEReader(Collections.<String, String>emptyMap());
    // This is a undesired single part state
    singlePartMIMEReader.setState(MultiPartMIMEReader.SingleReaderState.REQUESTED_DATA);
    reader.setCurrentSinglePartMIMEReader(singlePartMIMEReader);
    try {
        reader.registerReaderCallback(EMPTY_MULTI_PART_MIME_READER_CALLBACK);
        Assert.fail();
    } catch (StreamBusyException streamBusyException) {
    }
}
Also used : EntityStream(com.linkedin.r2.message.stream.entitystream.EntityStream) MultiPartReaderFinishedException(com.linkedin.multipart.exceptions.MultiPartReaderFinishedException) StreamBusyException(com.linkedin.multipart.exceptions.StreamBusyException) StreamRequest(com.linkedin.r2.message.stream.StreamRequest) Test(org.testng.annotations.Test)

Example 57 with EntityStream

use of com.linkedin.r2.message.stream.entitystream.EntityStream 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<>());
    _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<>();
        @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)

Example 58 with EntityStream

use of com.linkedin.r2.message.stream.entitystream.EntityStream 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<>();
    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 59 with EntityStream

use of com.linkedin.r2.message.stream.entitystream.EntityStream in project rest.li by linkedin.

the class TestUnstructuredDataClient method testClientWithStreamResponse.

@Test(dataProvider = "goodURLs")
public void testClientWithStreamResponse(String resourceURL) throws Throwable {
    ByteString expectedContent = ByteString.copy(GreetingUnstructuredDataUtils.UNSTRUCTURED_DATA_BYTES);
    URI testURI = URI.create(RestLiIntegrationTest.URI_PREFIX + resourceURL);
    EntityStream emptyStream = EntityStreams.emptyStream();
    StreamRequest streamRequest = new StreamRequestBuilder(testURI).build(emptyStream);
    final CountDownLatch latch = new CountDownLatch(1);
    final AtomicBoolean success = new AtomicBoolean();
    _client.streamRequest(streamRequest, new Callback<StreamResponse>() {

        @Override
        public void onError(Throwable e) {
            fail("failed to get response", e);
        }

        @Override
        public void onSuccess(StreamResponse result) {
            assertEquals(result.getHeader(RestConstants.HEADER_CONTENT_TYPE), GreetingUnstructuredDataUtils.MIME_TYPE);
            FullEntityReader fullEntityReader = new FullEntityReader(new Callback<ByteString>() {

                @Override
                public void onError(Throwable e) {
                    success.set(false);
                    latch.countDown();
                }

                @Override
                public void onSuccess(ByteString result) {
                    // Won't fail the test, only use to print out error
                    assertEquals(result, expectedContent);
                    // Will fail the test if content is not identical
                    success.set(result.equals(expectedContent));
                    latch.countDown();
                }
            });
            result.getEntityStream().setReader(fullEntityReader);
        }
    });
    latch.await(10, TimeUnit.SECONDS);
    if (!success.get())
        fail("Failed to read response data from stream!");
}
Also used : ByteString(com.linkedin.data.ByteString) StreamResponse(com.linkedin.r2.message.stream.StreamResponse) CountDownLatch(java.util.concurrent.CountDownLatch) URI(java.net.URI) StreamRequestBuilder(com.linkedin.r2.message.stream.StreamRequestBuilder) StreamRequest(com.linkedin.r2.message.stream.StreamRequest) EntityStream(com.linkedin.r2.message.stream.entitystream.EntityStream) FullEntityReader(com.linkedin.r2.message.stream.entitystream.FullEntityReader) AtomicBoolean(java.util.concurrent.atomic.AtomicBoolean) FutureCallback(com.linkedin.common.callback.FutureCallback) Callback(com.linkedin.common.callback.Callback) RestLiIntegrationTest(com.linkedin.restli.examples.RestLiIntegrationTest) Test(org.testng.annotations.Test)

Example 60 with EntityStream

use of com.linkedin.r2.message.stream.entitystream.EntityStream in project rest.li by linkedin.

the class TestRestLiServer method testUnstructuredDataGetWithBody.

@Test
public void testUnstructuredDataGetWithBody() throws Exception {
    EntityStream streams = EntityStreams.newEntityStream(new ByteStringWriter(ByteString.copy(new byte[] { 1, 2, 3, 4 })));
    Observer observer = new TestObserver();
    streams.addObserver(observer);
    final StreamRequest streamRequest = new StreamRequestBuilder(new URI("/reactiveGreetingCollectionUnstructuredData/hello")).setHeader(RestConstants.HEADER_RESTLI_PROTOCOL_VERSION, AllProtocolVersions.BASELINE_PROTOCOL_VERSION.toString()).build(streams);
    final Callback<StreamResponse> callback = new Callback<StreamResponse>() {

        @Override
        public void onSuccess(StreamResponse streamResponse) {
        }

        @Override
        public void onError(Throwable e) {
        }
    };
    _server.handleRequest(streamRequest, new RequestContext(), callback);
    assertTrue(((TestObserver) observer).isDone());
}
Also used : EntityStream(com.linkedin.r2.message.stream.entitystream.EntityStream) SinglePartMIMEFullReaderCallback(com.linkedin.multipart.utils.MIMETestUtils.SinglePartMIMEFullReaderCallback) Callback(com.linkedin.common.callback.Callback) MultiPartMIMEFullReaderCallback(com.linkedin.multipart.utils.MIMETestUtils.MultiPartMIMEFullReaderCallback) Observer(com.linkedin.r2.message.stream.entitystream.Observer) StreamResponse(com.linkedin.r2.message.stream.StreamResponse) FilterRequestContext(com.linkedin.restli.server.filter.FilterRequestContext) RequestContext(com.linkedin.r2.message.RequestContext) ByteStringWriter(com.linkedin.r2.message.stream.entitystream.ByteStringWriter) StreamRequestBuilder(com.linkedin.r2.message.stream.StreamRequestBuilder) URI(java.net.URI) StreamRequest(com.linkedin.r2.message.stream.StreamRequest) Test(org.testng.annotations.Test) AfterTest(org.testng.annotations.AfterTest) BeforeTest(org.testng.annotations.BeforeTest)

Aggregations

EntityStream (com.linkedin.r2.message.stream.entitystream.EntityStream)58 StreamResponse (com.linkedin.r2.message.stream.StreamResponse)29 Test (org.testng.annotations.Test)29 StreamRequest (com.linkedin.r2.message.stream.StreamRequest)27 CountDownLatch (java.util.concurrent.CountDownLatch)27 StreamRequestBuilder (com.linkedin.r2.message.stream.StreamRequestBuilder)22 ByteString (com.linkedin.data.ByteString)11 StreamResponseBuilder (com.linkedin.r2.message.stream.StreamResponseBuilder)10 FutureCallback (com.linkedin.common.callback.FutureCallback)9 AtomicInteger (java.util.concurrent.atomic.AtomicInteger)9 AbstractServiceTest (test.r2.integ.clientserver.providers.AbstractServiceTest)9 AtomicReference (java.util.concurrent.atomic.AtomicReference)8 Callback (com.linkedin.common.callback.Callback)7 Client (com.linkedin.r2.transport.common.Client)7 BytesWriter (test.r2.integ.helper.BytesWriter)7 ScheduledExecutorService (java.util.concurrent.ScheduledExecutorService)6 TimeoutException (java.util.concurrent.TimeoutException)6 StreamingCompressor (com.linkedin.r2.filter.compression.streaming.StreamingCompressor)5 None (com.linkedin.common.util.None)4 StreamEncodingType (com.linkedin.r2.filter.compression.streaming.StreamEncodingType)4