Examples with Writer com.linkedin.r2.message.stream.entitystream.Writer used on opensource projects

Example 46 with Writer

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

the class TestStreamingCompression method testDecompress.

private void testDecompress(StreamingCompressor compressor, byte[] uncompressed, byte[] compressed) throws CompressionException, ExecutionException, InterruptedException {
    ByteWriter writer = new ByteWriter(compressed);
    EntityStream compressedStream = EntityStreams.newEntityStream(writer);
    EntityStream uncompressedStream = compressor.inflate(compressedStream);
    FutureCallback<byte[]> callback = new FutureCallback<>();
    uncompressedStream.setReader(new ByteReader(callback));
    byte[] result = callback.get();
    Assert.assertEquals(result, uncompressed);
}

Example 47 with Writer

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

the class TestJetty404 method testJetty404.

// make sure jetty's default behavior will read all the request bytes in case of 404
// Known to be flaky in CI
@Test(retryAnalyzer = SingleRetry.class)
public void testJetty404() throws Exception {
    BytesWriter writer = new BytesWriter(200 * 1024, (byte) 100);
    final AtomicReference<Throwable> exRef = new AtomicReference<>();
    final CountDownLatch latch = new CountDownLatch(1);
    _client.streamRequest(new StreamRequestBuilder(Bootstrap.createHttpURI(PORT, URI.create("/wrong-path"))).build(EntityStreams.newEntityStream(writer)), new Callback<StreamResponse>() {

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

        @Override
        public void onSuccess(StreamResponse result) {
            latch.countDown();
        }
    });
    latch.await(5000, TimeUnit.MILLISECONDS);
    Assert.assertTrue(writer.isDone());
    Throwable ex = exRef.get();
    Assert.assertTrue(ex instanceof StreamException, "Expected StreamException but found: " + ex);
    StreamResponse response = ((StreamException) ex).getResponse();
    Assert.assertEquals(response.getStatus(), RestStatus.NOT_FOUND);
    response.getEntityStream().setReader(new DrainReader());
}

Example 48 with Writer

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

the class TestMIMEIntegrationReaderWriter method testMultipleBodies.

@Test(dataProvider = "chunkSizes")
public void testMultipleBodies(final int chunkSize) throws Exception {
    final MultiPartMIMEInputStream normalBodyInputStream = new MultiPartMIMEInputStream.Builder(new ByteArrayInputStream(_normalBody.getPartData().copyBytes()), scheduledExecutorService, _normalBody.getPartHeaders()).withWriteChunkSize(chunkSize).build();
    final MultiPartMIMEInputStream headerLessBodyInputStream = new MultiPartMIMEInputStream.Builder(new ByteArrayInputStream(_headerLessBody.getPartData().copyBytes()), scheduledExecutorService, _headerLessBody.getPartHeaders()).withWriteChunkSize(chunkSize).build();
    // Copying over empty ByteString, but let's keep things consistent.
    final MultiPartMIMEInputStream bodyLessBodyInputStream = new MultiPartMIMEInputStream.Builder(new ByteArrayInputStream(_bodyLessBody.getPartData().copyBytes()), scheduledExecutorService, _bodyLessBody.getPartHeaders()).withWriteChunkSize(chunkSize).build();
    final MultiPartMIMEInputStream purelyEmptyBodyInputStream = new MultiPartMIMEInputStream.Builder(new ByteArrayInputStream(_purelyEmptyBody.getPartData().copyBytes()), scheduledExecutorService, _purelyEmptyBody.getPartHeaders()).withWriteChunkSize(chunkSize).build();
    final MultiPartMIMEWriter writer = new MultiPartMIMEWriter.Builder("some preamble", "").appendDataSource(normalBodyInputStream).appendDataSource(headerLessBodyInputStream).appendDataSource(bodyLessBodyInputStream).appendDataSource(purelyEmptyBodyInputStream).build();
    executeRequestAndAssert(writer, Collections.unmodifiableList(Arrays.asList(_normalBody, _headerLessBody, _bodyLessBody, _purelyEmptyBody)));
}

Example 49 with Writer

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

the class TestMIMEChainingMultipleSources method testSinglePartDataSource.

// Test breakdown:
// 1. Main thread sends mime request
// 2. Server A sends a simple POST request to Server B
// 3. Server B sends back a mime response to Server A
// 4. Server A takes the original incoming request from the main thread + a local input
// stream + the first part from the incoming mime response from Server B.
// 5. Main thread then gets all of this and stores it.
// 6. Server A then drains and stores the rest of the parts from Server B's response.
@Test(dataProvider = "chunkSizes", groups = TESTNG_GROUP_KNOWN_ISSUE)
public void testSinglePartDataSource(final int chunkSize) throws Exception {
    _chunkSize = chunkSize;
    final List<MultiPartMIMEDataSourceWriter> dataSources = generateInputStreamDataSources(chunkSize, _scheduledExecutorService);
    final MultiPartMIMEWriter writer = new MultiPartMIMEWriter.Builder().appendDataSources(dataSources).build();
    final StreamRequest streamRequest = MultiPartMIMEStreamRequestFactory.generateMultiPartMIMEStreamRequest(Bootstrap.createHttpsURI(PORT_SERVER_A, SERVER_A_URI), "mixed", writer, Collections.<String, String>emptyMap());
    ClientMultiPartReceiver clientReceiver = new ClientMultiPartReceiver();
    Callback<StreamResponse> callback = generateSuccessChainCallback(clientReceiver);
    // Send the request to Server A
    _client.streamRequest(streamRequest, callback);
    _latch.await(TEST_TIMEOUT, TimeUnit.MILLISECONDS);
    // Verify client
    List<MIMETestUtils.SinglePartMIMEFullReaderCallback> clientSinglePartCallbacks = clientReceiver.getSinglePartMIMEReaderCallbacks();
    Assert.assertEquals(clientReceiver.getSinglePartMIMEReaderCallbacks().size(), 6);
    Assert.assertEquals(clientSinglePartCallbacks.get(0).getFinishedData(), BODY_1.getPartData());
    Assert.assertEquals(clientSinglePartCallbacks.get(0).getHeaders(), BODY_1.getPartHeaders());
    Assert.assertEquals(clientSinglePartCallbacks.get(1).getFinishedData(), BODY_5.getPartData());
    Assert.assertEquals(clientSinglePartCallbacks.get(1).getHeaders(), BODY_5.getPartHeaders());
    Assert.assertEquals(clientSinglePartCallbacks.get(2).getFinishedData(), BODY_A.getPartData());
    Assert.assertEquals(clientSinglePartCallbacks.get(2).getHeaders(), BODY_A.getPartHeaders());
    Assert.assertEquals(clientSinglePartCallbacks.get(3).getFinishedData(), BODY_B.getPartData());
    Assert.assertEquals(clientSinglePartCallbacks.get(3).getHeaders(), BODY_B.getPartHeaders());
    Assert.assertEquals(clientSinglePartCallbacks.get(4).getFinishedData(), BODY_C.getPartData());
    Assert.assertEquals(clientSinglePartCallbacks.get(4).getHeaders(), BODY_C.getPartHeaders());
    Assert.assertEquals(clientSinglePartCallbacks.get(5).getFinishedData(), BODY_D.getPartData());
    Assert.assertEquals(clientSinglePartCallbacks.get(5).getHeaders(), BODY_D.getPartHeaders());
    // Verify Server A
    List<MIMETestUtils.SinglePartMIMEFullReaderCallback> serverASinglePartCallbacks = _serverAMultiPartCallback.getSinglePartMIMEReaderCallbacks();
    Assert.assertEquals(serverASinglePartCallbacks.size(), 3);
    Assert.assertEquals(serverASinglePartCallbacks.get(0).getFinishedData(), BODY_2.getPartData());
    Assert.assertEquals(serverASinglePartCallbacks.get(0).getHeaders(), BODY_2.getPartHeaders());
    Assert.assertEquals(serverASinglePartCallbacks.get(1).getFinishedData(), BODY_3.getPartData());
    Assert.assertEquals(serverASinglePartCallbacks.get(1).getHeaders(), BODY_3.getPartHeaders());
    Assert.assertEquals(serverASinglePartCallbacks.get(2).getFinishedData(), BODY_4.getPartData());
    Assert.assertEquals(serverASinglePartCallbacks.get(2).getHeaders(), BODY_4.getPartHeaders());
}

Example 50 with Writer

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

the class TestMIMEChainingReader method testMimeReaderDataSource.

// Verifies that a multi part mime reader can be used as a data source to the writer.
// To make the test easier to write, we simply chain back to the client in the form of simulating a response.
@Test(dataProvider = "chunkSizes")
public void testMimeReaderDataSource(final int chunkSize) throws Exception {
    final List<MultiPartMIMEDataSourceWriter> dataSources = generateInputStreamDataSources(chunkSize, _scheduledExecutorService);
    final MultiPartMIMEWriter writer = new MultiPartMIMEWriter.Builder().appendDataSources(dataSources).build();
    final StreamRequest streamRequest = mock(StreamRequest.class);
    when(streamRequest.getEntityStream()).thenReturn(writer.getEntityStream());
    final String contentTypeHeader = "multipart/mixed; boundary=" + writer.getBoundary();
    when(streamRequest.getHeader(MultiPartMIMEUtils.CONTENT_TYPE_HEADER)).thenReturn(contentTypeHeader);
    // Client side preparation to read the part back on the callback
    // Note the chunks size will carry over since the client is controlling how much data he gets back
    // based on the chunk size he writes.
    final CountDownLatch latch = new CountDownLatch(1);
    ClientMultiPartMIMEReaderReceiverCallback _clientReceiver = new ClientMultiPartMIMEReaderReceiverCallback(latch);
    Callback<StreamResponse> callback = generateSuccessChainCallback(_clientReceiver);
    // Server side start
    MultiPartMIMEReader reader = MultiPartMIMEReader.createAndAcquireStream(streamRequest);
    ServerMultiPartMIMEChainReaderWriterCallback _serverSender = new ServerMultiPartMIMEChainReaderWriterCallback(callback, reader);
    reader.registerReaderCallback(_serverSender);
    latch.await(_testTimeout, TimeUnit.MILLISECONDS);
    // Verify client. No need to verify the server.
    List<MIMETestUtils.SinglePartMIMEFullReaderCallback> singlePartMIMEReaderCallbacks = _clientReceiver.getSinglePartMIMEReaderCallbacks();
    Assert.assertEquals(singlePartMIMEReaderCallbacks.size(), 4);
    Assert.assertEquals(singlePartMIMEReaderCallbacks.get(0).getFinishedData(), BODY_A.getPartData());
    Assert.assertEquals(singlePartMIMEReaderCallbacks.get(0).getHeaders(), BODY_A.getPartHeaders());
    Assert.assertEquals(singlePartMIMEReaderCallbacks.get(1).getFinishedData(), BODY_B.getPartData());
    Assert.assertEquals(singlePartMIMEReaderCallbacks.get(1).getHeaders(), BODY_B.getPartHeaders());
    Assert.assertEquals(singlePartMIMEReaderCallbacks.get(2).getFinishedData(), BODY_C.getPartData());
    Assert.assertEquals(singlePartMIMEReaderCallbacks.get(2).getHeaders(), BODY_C.getPartHeaders());
    Assert.assertEquals(singlePartMIMEReaderCallbacks.get(3).getFinishedData(), BODY_D.getPartData());
    Assert.assertEquals(singlePartMIMEReaderCallbacks.get(3).getHeaders(), BODY_D.getPartHeaders());
}