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

Example 76 with StreamRequestBuilder

use of com.linkedin.r2.message.stream.StreamRequestBuilder 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());
        }
    }
}

Example 77 with StreamRequestBuilder

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

the class TestHttp2NettyStreamClient method testRequestTimeout.

/**
 * Tests the condition that when a client request times out before the request is processed
 * by the server, the servlet implementation throws when attempting to read the request entity.
 */
@Test(enabled = false)
public void testRequestTimeout() throws Exception {
    final AtomicInteger serverIOExceptions = new AtomicInteger(0);
    final CountDownLatch exceptionLatch = new CountDownLatch(1);
    final CountDownLatch responseLatch = new CountDownLatch(1);
    final CountDownLatch serverLatch = new CountDownLatch(1);
    final HttpServerBuilder serverBuilder = new HttpServerBuilder();
    final Server server = serverBuilder.exceptionListener(throwable -> {
        if (throwable instanceof IOException) {
            serverIOExceptions.incrementAndGet();
            exceptionLatch.countDown();
        }
    }).responseLatch(serverLatch).build();
    final HttpClientBuilder clientBuilder = new HttpClientBuilder(_eventLoop, _scheduler);
    final Http2NettyStreamClient client = clientBuilder.setRequestTimeout(500).buildHttp2StreamClient();
    try {
        server.start();
        // Sends the stream request
        final StreamRequestBuilder builder = new StreamRequestBuilder(new URI(URL));
        final StreamRequest request = builder.setMethod(METHOD).setHeader(HttpHeaderNames.HOST.toString(), HOST_NAME.toString()).build(EntityStreams.newEntityStream(new ByteStringWriter(ByteString.copy(new byte[REQUEST_SIZE]))));
        client.streamRequest(request, new RequestContext(), new HashMap<>(), response -> responseLatch.countDown());
        // Waits for request to timeout
        Thread.sleep(1000);
        // Allows server to process request
        serverLatch.countDown();
    } finally {
        if (!responseLatch.await(TEST_TIMEOUT, TimeUnit.MILLISECONDS)) {
            Assert.fail("Timeout waiting for response latch");
        }
        if (!exceptionLatch.await(TEST_TIMEOUT, TimeUnit.MILLISECONDS)) {
            Assert.fail("Timeout waiting for exception latch");
        }
        server.stop();
    }
    // Expects two IOExceptions thrown by the server. One for the initial OPTIONS upgrade request and one for
    // the actual GET request.
    Assert.assertEquals(serverIOExceptions.get(), 2);
}

Example 78 with StreamRequestBuilder

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

the class StreamRequestGenerator method nextMessage.

@Override
public StreamRequest nextMessage() {
    if (_msgCounter.getAndDecrement() > 0) {
        StreamRequestBuilder builder = new StreamRequestBuilder(_uri);
        builder.setMethod(HTTP_POST_METHOD);
        for (int i = 0; i < _numHeaders; i++) {
            builder.setHeader(STATIC_HEADER_PREFIX + i, _headerContent);
        }
        return builder.build(EntityStreams.newEntityStream(new PerfStreamWriter(_msgSize)));
    } else {
        return null;
    }
}

Example 79 with StreamRequestBuilder

use of com.linkedin.r2.message.stream.StreamRequestBuilder 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!");
}

Example 80 with StreamRequestBuilder

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

the class TestRestLiServer method testValidUnstructuredDataRequest.

@Test(dataProvider = "validClientProtocolVersionDataStreamOnly")
public void testValidUnstructuredDataRequest(RestLiServer server, ProtocolVersion clientProtocolVersion, String headerConstant) throws URISyntaxException, IOException {
    StreamRequest streamRequest = new StreamRequestBuilder(new URI("/feedDownloads/1")).setHeader(headerConstant, clientProtocolVersion.toString()).build(EntityStreams.emptyStream());
    final FeedDownloadResource resource = getMockResource(FeedDownloadResource.class);
    resource.get(eq(1L), anyObject(UnstructuredDataWriter.class));
    EasyMock.expectLastCall().andDelegateTo(new FeedDownloadResource()).once();
    replay(resource);
    @SuppressWarnings("unchecked") Callback<StreamResponse> r2Callback = createMock(Callback.class);
    final Capture<StreamResponse> streamResponse = EasyMock.newCapture();
    r2Callback.onSuccess(capture(streamResponse));
    expectLastCall().once();
    replay(r2Callback);
    RequestContext requestContext = new RequestContext();
    server.handleRequest(streamRequest, requestContext, r2Callback);
    verify(resource);
    verify(r2Callback);
    assertNotNull(streamResponse);
    assertEquals(streamResponse.getValue().getHeader(RestConstants.HEADER_CONTENT_TYPE), FeedDownloadResource.CONTENT_TYPE);
    FullEntityReader fullEntityReader = new FullEntityReader(new Callback<ByteString>() {

        @Override
        public void onError(Throwable e) {
            fail("Error inside callback!! Failed to read response data from stream!", e);
        }

        @Override
        public void onSuccess(ByteString result) {
            assertEquals(result.copyBytes(), FeedDownloadResource.CONTENT);
        }
    });
    streamResponse.getValue().getEntityStream().setReader(fullEntityReader);
}