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

Example 81 with StreamResponse

use of com.linkedin.r2.message.stream.StreamResponse 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 82 with StreamResponse

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

Example 83 with StreamResponse

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

the class TestHttpBridge method testStreamToHttpErrorMessage.

@Test
public void testStreamToHttpErrorMessage() throws TimeoutException, InterruptedException {
    URI uri = URI.create("http://some.host/thisShouldAppearInTheErrorMessage");
    RestRequest r = new RestRequestBuilder(uri).build();
    FutureCallback<StreamResponse> futureCallback = new FutureCallback<StreamResponse>();
    TransportCallback<StreamResponse> callback = new TransportCallbackAdapter<StreamResponse>(futureCallback);
    TransportCallback<StreamResponse> bridgeCallback = HttpBridge.streamToHttpCallback(callback, Messages.toStreamRequest(r));
    bridgeCallback.onResponse(TransportResponseImpl.<StreamResponse>error(new Exception()));
    try {
        futureCallback.get(30, TimeUnit.SECONDS);
        Assert.fail("get should have thrown exception");
    } catch (ExecutionException e) {
        Assert.assertTrue(e.getCause().getMessage().contains(uri.toString()));
    }
}

Example 84 with StreamResponse

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

the class TestDisruptFilter method testStreamTimeoutDisrupt.

@Test
public void testStreamTimeoutDisrupt() throws Exception {
    final RequestContext requestContext = new RequestContext();
    requestContext.putLocalAttr(DISRUPT_CONTEXT_KEY, DisruptContexts.timeout());
    final DisruptFilter filter = new DisruptFilter(_scheduler, _executor, REQUEST_TIMEOUT);
    final CountDownLatch latch = new CountDownLatch(1);
    final AtomicBoolean success = new AtomicBoolean(false);
    final NextFilter<StreamRequest, StreamResponse> next = new NextFilter<StreamRequest, StreamResponse>() {

        @Override
        public void onRequest(StreamRequest restRequest, RequestContext requestContext, Map<String, String> wireAttrs) {
            latch.countDown();
        }

        @Override
        public void onResponse(StreamResponse restResponse, RequestContext requestContext, Map<String, String> wireAttrs) {
            latch.countDown();
        }

        @Override
        public void onError(Throwable ex, RequestContext requestContext, Map<String, String> wireAttrs) {
            success.set(ex instanceof TimeoutException);
            latch.countDown();
        }
    };
    filter.onStreamRequest(new StreamRequestBuilder(new URI(URI)).build(EntityStreams.emptyStream()), requestContext, Collections.emptyMap(), next);
    Assert.assertTrue(latch.await(TEST_TIMEOUT, TimeUnit.MILLISECONDS), "Missing NextFilter invocation");
    Assert.assertTrue(success.get(), "Unexpected method invocation");
}

Example 85 with StreamResponse

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

the class TestStreamFilterAdapters method testResponseFilterAdapterPassThrough.

@Test
public void testResponseFilterAdapterPassThrough() {
    FilterChain fc = adaptAndCreateFilterChain(new RestFilter() {

        @Override
        public void onRestResponse(RestResponse res, RequestContext requestContext, Map<String, String> wireAttrs, NextFilter<RestRequest, RestResponse> nextFilter) {
            nextFilter.onResponse(res, requestContext, wireAttrs);
        }

        @Override
        public void onRestError(Throwable ex, RequestContext requestContext, Map<String, String> wireAttrs, NextFilter<RestRequest, RestResponse> nextFilter) {
            nextFilter.onError(ex, requestContext, wireAttrs);
        }
    });
    fc.onStreamResponse(simpleStreamResponse("12345"), FilterUtil.emptyRequestContext(), FilterUtil.emptyWireAttrs());
    StreamResponse capturedResponse = _beforeFilter.getResponse();
    capturedResponse.getEntityStream().setReader(new FullEntityReader(new Callback<ByteString>() {

        @Override
        public void onError(Throwable e) {
            Assert.fail("should not happen");
        }

        @Override
        public void onSuccess(ByteString result) {
            Assert.assertEquals(result.asString("UTF8"), "12345");
        }
    }));
    fc.onStreamError(simpleStreamException("12345"), FilterUtil.emptyRequestContext(), FilterUtil.emptyWireAttrs());
    Throwable capturedEx = _beforeFilter.getThrowable();
    Assert.assertTrue(capturedEx instanceof StreamException);
    ((StreamException) capturedEx).getResponse().getEntityStream().setReader(new FullEntityReader(new Callback<ByteString>() {

        @Override
        public void onError(Throwable e) {
            Assert.fail("should not happen");
        }

        @Override
        public void onSuccess(ByteString result) {
            Assert.assertEquals(result.asString("UTF8"), "12345");
        }
    }));
}